


import os
import sys
import warnings
from dataclasses import dataclass, field
from io import BytesIO
from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Optional, Union

import numpy as np
import pyarrow as pa

from .. import config
from ..download.streaming_download_manager import xopen
from ..table import array_cast
from ..utils.file_utils import is_local_path
from ..utils.py_utils import first_non_null_value, no_op_if_value_is_null, string_to_dict


if TYPE_CHECKING:
				import PIL.Image

				from .features import FeatureType


_lowercase							:   Tuple						=						None
_lowercase							:   Union[str, Any]						=						"<" if sys.byteorder == "little" else ">"
# Origin: https://github.com/python-pillow/Pillow/blob/698951e19e19972aeed56df686868f1329981c12/src/PIL/Image.py#L3126 minus "|i1" which values are not preserved correctly when saving and loading an image
_lowercase							:   Union[str, Any]						=						[
    np.dtype("|b1"),
    np.dtype("|u1"),
    np.dtype("<u2"),
    np.dtype(">u2"),
    np.dtype("<i2"),
    np.dtype(">i2"),
    np.dtype("<u4"),
    np.dtype(">u4"),
    np.dtype("<i4"),
    np.dtype(">i4"),
    np.dtype("<f4"),
    np.dtype(">f4"),
    np.dtype("<f8"),
    np.dtype(">f8"),
]






@dataclass
class    _UpperCamelCase     :

				"""simple docstring"""







				lowerCAmelCase					=			True
				lowerCAmelCase					=			None
				# Automatically constructed
				lowerCAmelCase					=			'PIL.Image.Image'
				lowerCAmelCase					=			pa.struct({'bytes': pa.binary(), 'path': pa.string()}					)
				lowerCAmelCase					=			field(default='Image'     ,		init=lowerCAmelCase__     ,		repr=lowerCAmelCase__					)

				def __call__(    self )						->					List[Any]:
					return self.pa_type

				def   _UpperCAmelCase       (    self					, a__ )						->					str:

					if config.PIL_AVAILABLE:
						import PIL.Image
					else:
						raise ImportError("""To support encoding images, please install 'Pillow'.""" )

					if isinstance(_lowerCamelCase					, _lowerCamelCase ):
						A             =					np.array(_lowerCamelCase )

					if isinstance(_lowerCamelCase					, _lowerCamelCase ):
						return {"path": value, "bytes": None}
					elif isinstance(_lowerCamelCase					, _lowerCamelCase ):
						return {"path": None, "bytes": value}
					elif isinstance(_lowerCamelCase					, np.ndarray ):
						# convert the image array to PNG/TIFF bytes
						return encode_np_array(_lowerCamelCase )
					elif isinstance(_lowerCamelCase					, PIL.Image.Image ):
						# convert the PIL image to bytes (default format is PNG/TIFF)
						return encode_pil_image(_lowerCamelCase )
					elif value.get("""path""" ) is not None and os.path.isfile(value["""path"""] ):
						# we set "bytes": None to not duplicate the data if they're already available locally
						return {"bytes": None, "path": value.get("""path""" )}
					elif value.get("""bytes""" ) is not None or value.get("""path""" ) is not None:
						# store the image bytes, and path is used to infer the image format using the file extension
						return {"bytes": value.get("""bytes""" ), "path": value.get("""path""" )}
					else:
						raise ValueError(
						    f'An image sample should have one of \'path\' or \'bytes\' but they are missing or None in {value}.' )

				def   _UpperCAmelCase       (    self					, a__					, a__=None )						->					int:

					if not self.decode:
						raise RuntimeError("""Decoding is disabled for this feature. Please use Image(decode=True) instead.""" )

					if config.PIL_AVAILABLE:
						import PIL.Image
					else:
						raise ImportError("""To support decoding images, please install 'Pillow'.""" )

					if token_per_repo_id is None:
						A             =					{}

					A             =					value["""path"""], value["""bytes"""]
					if bytes_ is None:
						if path is None:
							raise ValueError(f'An image should have one of \'path\' or \'bytes\' but both are None in {value}.' )
						else:
							if is_local_path(_lowerCamelCase ):
								A             =					PIL.Image.open(_lowerCamelCase )
							else:
								A             =					path.split("""::""" )[-1]
								try:
									A             =					string_to_dict(_lowerCamelCase					, config.HUB_DATASETS_URL )["""repo_id"""]
									A             =					token_per_repo_id.get(_lowerCamelCase )
								except ValueError:
									A             =					None
								with xopen(_lowerCamelCase					, """rb"""					, use_auth_token=_lowerCamelCase ) as f:
									A             =					BytesIO(f.read() )
								A             =					PIL.Image.open(bytes_ )
					else:
						A             =					PIL.Image.open(BytesIO(bytes_ ) )
					image.load()  # to avoid "Too many open files" errors
					return image

				def   _UpperCAmelCase       (    self )						->					List[Any]:
					from .features import Value

					return (
					    self
					    if self.decode
					    else {
					        "bytes": Value("""binary""" ),
					        "path": Value("""string""" ),
					    }
					)

				def   _UpperCAmelCase       (    self					, a__ )						->					Tuple:

					if pa.types.is_string(storage.type ):
						A             =					pa.array([None] * len(_lowerCamelCase )					, type=pa.binary() )
						A             =					pa.StructArray.from_arrays([bytes_array, storage]					, ["""bytes""", """path"""]					, mask=storage.is_null() )
					elif pa.types.is_binary(storage.type ):
						A             =					pa.array([None] * len(_lowerCamelCase )					, type=pa.string() )
						A             =					pa.StructArray.from_arrays([storage, path_array]					, ["""bytes""", """path"""]					, mask=storage.is_null() )
					elif pa.types.is_struct(storage.type ):
						if storage.type.get_field_index("""bytes""" ) >= 0:
							A             =					storage.field("""bytes""" )
						else:
							A             =					pa.array([None] * len(_lowerCamelCase )					, type=pa.binary() )
						if storage.type.get_field_index("""path""" ) >= 0:
							A             =					storage.field("""path""" )
						else:
							A             =					pa.array([None] * len(_lowerCamelCase )					, type=pa.string() )
						A             =					pa.StructArray.from_arrays([bytes_array, path_array]					, ["""bytes""", """path"""]					, mask=storage.is_null() )
					elif pa.types.is_list(storage.type ):
						A             =					pa.array(
						    [encode_np_array(np.array(_lowerCamelCase ) )["""bytes"""] if arr is not None else None for arr in storage.to_pylist()]					, type=pa.binary()					, )
						A             =					pa.array([None] * len(_lowerCamelCase )					, type=pa.string() )
						A             =					pa.StructArray.from_arrays(
						    [bytes_array, path_array]					, ["""bytes""", """path"""]					, mask=bytes_array.is_null() )
					return array_cast(_lowerCamelCase					, self.pa_type )

				def   _UpperCAmelCase       (    self					, a__ )						->					Any:

					@no_op_if_value_is_null
					def path_to_bytes(a__ ):
						with xopen(_lowerCamelCase					, """rb""" ) as f:
							A             =					f.read()
						return bytes_

					A             =					pa.array(
					    [
					        (path_to_bytes(x["""path"""] ) if x["""bytes"""] is None else x["""bytes"""]) if x is not None else None
					        for x in storage.to_pylist()
					    ]					, type=pa.binary()					, )
					A             =					pa.array(
					    [os.path.basename(_lowerCamelCase ) if path is not None else None for path in storage.field("""path""" ).to_pylist()]					, type=pa.string()					, )
					A             =					pa.StructArray.from_arrays([bytes_array, path_array]					, ["""bytes""", """path"""]					, mask=bytes_array.is_null() )
					return array_cast(_lowerCamelCase					, self.pa_type )





def 	_lowerCAmelCase					( )			-> List[str]:





	"""simple docstring"""







	if config.PIL_AVAILABLE:
		import PIL.Image
	else:
		raise ImportError("""To support encoding images, please install 'Pillow'."""    )

	global _IMAGE_COMPRESSION_FORMATS
	if _IMAGE_COMPRESSION_FORMATS is None:
		PIL.Image.init()
		A             =					list(set(PIL.Image.OPEN.keys()    ) & set(PIL.Image.SAVE.keys()    )    )
	return _IMAGE_COMPRESSION_FORMATS





def 	_lowerCAmelCase					( UpperCamelCase__:  List[str]    )			-> bytes:





	"""simple docstring"""







	A             =					BytesIO()
	if image.format in list_image_compression_formats():
		A             =					image.format
	else:
		A             =					"""PNG""" if image.mode in ["""1""", """L""", """LA""", """RGB""", """RGBA"""] else """TIFF"""
	image.save(UpperCamelCase__				,       format=UpperCamelCase__    )
	return buffer.getvalue()





def 	_lowerCAmelCase					( UpperCamelCase__:  Dict    )			-> dict:





	"""simple docstring"""







	if hasattr(UpperCamelCase__				,       """filename"""    ) and image.filename != "":
		return {"path": image.filename, "bytes": None}
	else:
		return {"path": None, "bytes": image_to_bytes(UpperCamelCase__    )}





def 	_lowerCAmelCase					( UpperCamelCase__:  Any    )			-> dict:





	"""simple docstring"""







	if config.PIL_AVAILABLE:
		import PIL.Image
	else:
		raise ImportError("""To support encoding images, please install 'Pillow'."""    )

	A             =					array.dtype
	A             =					dtype.byteorder if dtype.byteorder != """=""" else _NATIVE_BYTEORDER
	A             =					dtype.kind
	A             =					dtype.itemsize

	A             =					None

	# Multi-channel array case (only np.dtype("|u1") is allowed)
	if array.shape[2:]:
		A             =					np.dtype("""|u1"""    )
		if dtype_kind not in ["u", "i"]:
			raise TypeError(
			    f'Unsupported array dtype {dtype} for image encoding. Only {dest_dtype} is supported for multi-channel arrays.'    )
		if dtype is not dest_dtype:
			warnings.warn(f'Downcasting array dtype {dtype} to {dest_dtype} to be compatible with \'Pillow\''    )
    # Exact match
	elif dtype in _VALID_IMAGE_ARRAY_DTPYES:
		A             =					dtype
	else:  # Downcast the type within the kind (np.can_cast(from_type, to_type, casting="same_kind") doesn't behave as expected, so do it manually)
		while dtype_itemsize >= 1:
			A             =					dtype_byteorder + dtype_kind + str(UpperCamelCase__    )
			A             =					np.dtype(UpperCamelCase__    )
			if dest_dtype in _VALID_IMAGE_ARRAY_DTPYES:
				warnings.warn(f'Downcasting array dtype {dtype} to {dest_dtype} to be compatible with \'Pillow\''    )
				break
			else:
				dtype_itemsize //= 2

	if dest_dtype is None:
		raise TypeError(
		    f'Cannot convert dtype {dtype} to a valid image dtype. Valid image dtypes: {_VALID_IMAGE_ARRAY_DTPYES}'    )
	A             =					PIL.Image.fromarray(array.astype(UpperCamelCase__    )    )
	return {"path": None, "bytes": image_to_bytes(UpperCamelCase__    )}





def 	_lowerCAmelCase					( UpperCamelCase__:  Any    )			-> List[dict]:





	"""simple docstring"""







	if config.PIL_AVAILABLE:
		import PIL.Image
	else:
		raise ImportError("""To support encoding images, please install 'Pillow'."""    )

	if objs:
		A             =					first_non_null_value(UpperCamelCase__    )
		if isinstance(UpperCamelCase__				,       UpperCamelCase__    ):
			return [{"path": obj, "bytes": None} if obj is not None else None for obj in objs]
		if isinstance(UpperCamelCase__				,       np.ndarray    ):
			A             =					no_op_if_value_is_null(UpperCamelCase__    )
			return [obj_to_image_dict_func(UpperCamelCase__    ) for obj in objs]
		elif isinstance(UpperCamelCase__				,       PIL.Image.Image    ):
			A             =					no_op_if_value_is_null(UpperCamelCase__    )
			return [obj_to_image_dict_func(UpperCamelCase__    ) for obj in objs]
		else:
			return objs
	else:
		return objs






"""simple docstring"""
import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyInpaintPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()

class       lowerCamelCase	(       lowerCAmelCase__					,					unittest.TestCase   ):


  '''simple docstring'''







  SCREAMING_SNAKE_CASE					      =  KandinskyInpaintPipeline
  SCREAMING_SNAKE_CASE					      =  ['prompt', 'image_embeds', 'negative_image_embeds', 'image', 'mask_image']
  SCREAMING_SNAKE_CASE					      =  [
      'prompt',
      'negative_prompt',
      'image_embeds',
      'negative_image_embeds',
      'image',
      'mask_image',
  ]
  SCREAMING_SNAKE_CASE					      =  [
      'generator',
      'height',
      'width',
      'latents',
      'guidance_scale',
      'negative_prompt',
      'num_inference_steps',
      'return_dict',
      'guidance_scale',
      'num_images_per_prompt',
      'output_type',
      'return_dict',
  ]
  SCREAMING_SNAKE_CASE					      =  False



  @property
  def  _a      (self	):


        """simple docstring"""

        return 32



  @property
  def  _a      (self	):


        """simple docstring"""

        return 32



  @property
  def  _a      (self	):


        """simple docstring"""

        return self.time_input_dim



  @property
  def  _a      (self	):


        """simple docstring"""

        return self.time_input_dim * 4



  @property
  def  _a      (self	):


        """simple docstring"""

        return 100



  @property
  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Optional[int]							=    XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base"""	)
        return tokenizer



  @property
  def  _a      (self	):


        """simple docstring"""

        torch.manual_seed(0	)
        UpperCAmelCase__			:							Optional[int]							=    MCLIPConfig(
            numDims=self.cross_attention_dim							,       transformerDimensions=self.text_embedder_hidden_size							,       hidden_size=self.text_embedder_hidden_size							,       intermediate_size=37							,       num_attention_heads=4							,       num_hidden_layers=5							,       vocab_size=1005							,       )

        UpperCAmelCase__			:							List[Any]							=    MultilingualCLIP(_lowerCamelCase	)
        UpperCAmelCase__			:							Optional[Any]							=    text_encoder.eval()

        return text_encoder



  @property
  def  _a      (self	):


        """simple docstring"""

        torch.manual_seed(0	)

        UpperCAmelCase__			:							Dict							=    {
            """in_channels""": 9,
            # Out channels is double in channels because predicts mean and variance
            """out_channels""": 8,
            """addition_embed_type""": """text_image""",
            """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
            """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
            """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
            """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
            """layers_per_block""": 1,
            """encoder_hid_dim""": self.text_embedder_hidden_size,
            """encoder_hid_dim_type""": """text_image_proj""",
            """cross_attention_dim""": self.cross_attention_dim,
            """attention_head_dim""": 4,
            """resnet_time_scale_shift""": """scale_shift""",
            """class_embed_type""": None,
        }

        UpperCAmelCase__			:							List[Any]							=    UNetaDConditionModel(**_lowerCamelCase	)
        return model



  @property
  def  _a      (self	):


        """simple docstring"""

        return {
            "block_out_channels": [32, 64],
            "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
            "in_channels": 3,
            "latent_channels": 4,
            "layers_per_block": 1,
            "norm_num_groups": 8,
            "norm_type": "spatial",
            "num_vq_embeddings": 12,
            "out_channels": 3,
            "up_block_types": [
                "AttnUpDecoderBlock2D",
                "UpDecoderBlock2D",
            ],
            "vq_embed_dim": 4,
        }



  @property
  def  _a      (self	):


        """simple docstring"""

        torch.manual_seed(0	)
        UpperCAmelCase__			:							Any							=    VQModel(**self.dummy_movq_kwargs	)
        return model



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Dict							=    self.dummy_text_encoder
        UpperCAmelCase__			:							Optional[Any]							=    self.dummy_tokenizer
        UpperCAmelCase__			:							Optional[Any]							=    self.dummy_unet
        UpperCAmelCase__			:							List[Any]							=    self.dummy_movq

        UpperCAmelCase__			:							Union[str, Any]							=    DDIMScheduler(
            num_train_timesteps=1000							,       beta_schedule="""linear"""							,       beta_start=0.00_085							,       beta_end=0.012							,       clip_sample=_lowerCamelCase							,       set_alpha_to_one=_lowerCamelCase							,       steps_offset=1							,       prediction_type="""epsilon"""							,       thresholding=_lowerCamelCase							,       )

        UpperCAmelCase__			:							Any							=    {
            """text_encoder""": text_encoder,
            """tokenizer""": tokenizer,
            """unet""": unet,
            """scheduler""": scheduler,
            """movq""": movq,
        }

        return components



  def  _a      (self							,       _lowerCamelCase							,       _lowerCamelCase=0	):


        """simple docstring"""

        UpperCAmelCase__			:							Dict							=    floats_tensor((1, self.cross_attention_dim)							,       rng=random.Random(_lowerCamelCase	)	).to(_lowerCamelCase	)
        UpperCAmelCase__			:							Optional[int]							=    floats_tensor((1, self.cross_attention_dim)							,       rng=random.Random(seed + 1	)	).to(_lowerCamelCase	)
        # create init_image
        UpperCAmelCase__			:							Optional[int]							=    floats_tensor((1, 3, 64, 64)							,       rng=random.Random(_lowerCamelCase	)	).to(_lowerCamelCase	)
        UpperCAmelCase__			:							Optional[int]							=    image.cpu().permute(0							,       2							,       3							,       1	)[0]
        UpperCAmelCase__			:							Optional[Any]							=    Image.fromarray(np.uinta(_lowerCamelCase	)	).convert("""RGB"""	).resize((256, 256)	)
        # create mask
        UpperCAmelCase__			:							List[Any]							=    np.ones((64, 64)							,       dtype=np.floataa	)
        UpperCAmelCase__			:							Optional[int]							=    0

        if str(_lowerCamelCase	).startswith("""mps"""	):
              UpperCAmelCase__			:							Tuple							=    torch.manual_seed(_lowerCamelCase	)
        else:
              UpperCAmelCase__			:							Optional[int]							=    torch.Generator(device=_lowerCamelCase	).manual_seed(_lowerCamelCase	)
        UpperCAmelCase__			:							Any							=    {
            """prompt""": """horse""",
            """image""": init_image,
            """mask_image""": mask,
            """image_embeds""": image_embeds,
            """negative_image_embeds""": negative_image_embeds,
            """generator""": generator,
            """height""": 64,
            """width""": 64,
            """num_inference_steps""": 2,
            """guidance_scale""": 4.0,
            """output_type""": """np""",
        }
        return inputs



  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							Union[str, Any]							=    """cpu"""

        UpperCAmelCase__			:							Optional[Any]							=    self.get_dummy_components()

        UpperCAmelCase__			:							Optional[Any]							=    self.pipeline_class(**_lowerCamelCase	)
        UpperCAmelCase__			:							Union[str, Any]							=    pipe.to(_lowerCamelCase	)

        pipe.set_progress_bar_config(disable=_lowerCamelCase	)

        UpperCAmelCase__			:							List[str]							=    pipe(**self.get_dummy_inputs(_lowerCamelCase	)	)
        UpperCAmelCase__			:							Optional[int]							=    output.images

        UpperCAmelCase__			:							Union[str, Any]							=    pipe(
            **self.get_dummy_inputs(_lowerCamelCase	)							,       return_dict=_lowerCamelCase							,       )[0]

        UpperCAmelCase__			:							Union[str, Any]							=    image[0, -3:, -3:, -1]
        UpperCAmelCase__			:							Optional[Any]							=    image_from_tuple[0, -3:, -3:, -1]

        print(F"""image.shape {image.shape}"""	)

        assert image.shape == (1, 64, 64, 3)

        UpperCAmelCase__			:							Any							=    np.array(
            [0.8_326_919, 0.73_790_467, 0.20_918_581, 0.9_309_612, 0.5_511_791, 0.43_713_328, 0.5_513_321, 0.49_922_934, 0.59_497_786]	)

        assert (
            np.abs(image_slice.flatten() - expected_slice	).max() < 1e-2
        ), F""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
        assert (
            np.abs(image_from_tuple_slice.flatten() - expected_slice	).max() < 1e-2
        ), F""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""




  def  _a      (self	):


        """simple docstring"""

        super().test_inference_batch_single_identical(expected_max_diff=3e-3	)




@slow
@require_torch_gpu
class       lowerCamelCase	(       unittest.TestCase   ):


  '''simple docstring'''







  def  _a      (self	):


        """simple docstring"""

        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()




  def  _a      (self	):


        """simple docstring"""

        UpperCAmelCase__			:							str							=    load_numpy(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
            """/kandinsky/kandinsky_inpaint_cat_with_hat_fp16.npy"""	)

        UpperCAmelCase__			:							Union[str, Any]							=    load_image(
            """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/kandinsky/cat.png"""	)
        UpperCAmelCase__			:							List[str]							=    np.ones((768, 768)							,       dtype=np.floataa	)
        UpperCAmelCase__			:							Optional[int]							=    0

        UpperCAmelCase__			:							int							=    """a hat"""

        UpperCAmelCase__			:							List[Any]							=    KandinskyPriorPipeline.from_pretrained(
            """kandinsky-community/kandinsky-2-1-prior"""							,       torch_dtype=torch.floataa	)
        pipe_prior.to(_lowerCamelCase	)

        UpperCAmelCase__			:							str							=    KandinskyInpaintPipeline.from_pretrained(
            """kandinsky-community/kandinsky-2-1-inpaint"""							,       torch_dtype=torch.floataa	)
        UpperCAmelCase__			:							int							=    pipeline.to(_lowerCamelCase	)
        pipeline.set_progress_bar_config(disable=_lowerCamelCase	)

        UpperCAmelCase__			:							int							=    torch.Generator(device="""cpu"""	).manual_seed(0	)
        UpperCAmelCase__					,		UpperCAmelCase__			:							List[str]							=    pipe_prior(
            _lowerCamelCase							,       generator=_lowerCamelCase							,       num_inference_steps=5							,       negative_prompt=""""""							,       ).to_tuple()

        UpperCAmelCase__			:							Union[str, Any]							=    pipeline(
            _lowerCamelCase							,       image=_lowerCamelCase							,       mask_image=_lowerCamelCase							,       image_embeds=_lowerCamelCase							,       negative_image_embeds=_lowerCamelCase							,       generator=_lowerCamelCase							,       num_inference_steps=100							,       height=768							,       width=768							,       output_type="""np"""							,       )

        UpperCAmelCase__			:							Any							=    output.images[0]

        assert image.shape == (768, 768, 3)

        assert_mean_pixel_difference(_lowerCamelCase							,       _lowerCamelCase	)

