





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


__A				:    Optional[List[str]]       =   None
__A				:    int       =   "<" 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
__A				:    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 	A_       :
	UpperCAmelCase__					=			True
	UpperCAmelCase__					=			None
	# Automatically constructed
	UpperCAmelCase__					=			"PIL.Image.Image"
	UpperCAmelCase__					=			pa.struct({'''bytes''': pa.binary(), '''path''': pa.string()}			)
	UpperCAmelCase__					=			field(default='''Image'''     ,							init=a_     ,							repr=a_			)
	def __call__(    self		):






				'''simple docstring'''







				return self.pa_type
	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







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

				if isinstance(_A							,					_A		):
							UpperCAmelCase   		=				np.array(_A		)

				if isinstance(_A							,					_A		):
							return {"path": value, "bytes": None}
				elif isinstance(_A							,					_A		):
							return {"path": None, "bytes": value}
				elif isinstance(_A							,					np.ndarray		):
							# convert the image array to PNG/TIFF bytes
							return encode_np_array(_A		)
				elif isinstance(_A							,					PIL.Image.Image		):
							# convert the PIL image to bytes (default format is PNG/TIFF)
							return encode_pil_image(_A		)
				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      _lowercase				(    self							,					_A							,					_A=None		):






				'''simple docstring'''







				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:
							UpperCAmelCase   		=				{}

				UpperCAmelCase		,			UpperCAmelCase   		=				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(_A		):
													UpperCAmelCase   		=				PIL.Image.open(_A		)
										else:
													UpperCAmelCase   		=				path.split('''::'''		)[-1]
													try:
																UpperCAmelCase   		=				string_to_dict(_A							,					config.HUB_DATASETS_URL		)['''repo_id''']
																UpperCAmelCase   		=				token_per_repo_id.get(_A		)
													except ValueError:
																UpperCAmelCase   		=				None
													with xopen(_A							,					'''rb'''							,					use_auth_token=_A		) as f:
																UpperCAmelCase   		=				BytesIO(f.read()		)
													UpperCAmelCase   		=				PIL.Image.open(bytes_		)
				else:
							UpperCAmelCase   		=				PIL.Image.open(BytesIO(bytes_		)		)
				image.load()  # to avoid "Too many open files" errors
				return image
	def      _lowercase				(    self		):






				'''simple docstring'''







				from .features import Value

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






				'''simple docstring'''







				if pa.types.is_string(storage.type		):
							UpperCAmelCase   		=				pa.array([None] * len(_A		)							,					type=pa.binary()		)
							UpperCAmelCase   		=				pa.StructArray.from_arrays([bytes_array, storage]							,					['''bytes''', '''path''']							,					mask=storage.is_null()		)
				elif pa.types.is_binary(storage.type		):
							UpperCAmelCase   		=				pa.array([None] * len(_A		)							,					type=pa.string()		)
							UpperCAmelCase   		=				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:
										UpperCAmelCase   		=				storage.field('''bytes'''		)
							else:
										UpperCAmelCase   		=				pa.array([None] * len(_A		)							,					type=pa.binary()		)
							if storage.type.get_field_index('''path'''		) >= 0:
										UpperCAmelCase   		=				storage.field('''path'''		)
							else:
										UpperCAmelCase   		=				pa.array([None] * len(_A		)							,					type=pa.string()		)
							UpperCAmelCase   		=				pa.StructArray.from_arrays([bytes_array, path_array]							,					['''bytes''', '''path''']							,					mask=storage.is_null()		)
				elif pa.types.is_list(storage.type		):
							UpperCAmelCase   		=				pa.array(
							    [encode_np_array(np.array(_A		)		)['''bytes'''] if arr is not None else None for arr in storage.to_pylist()]							,					type=pa.binary()							,					)
							UpperCAmelCase   		=				pa.array([None] * len(_A		)							,					type=pa.string()		)
							UpperCAmelCase   		=				pa.StructArray.from_arrays(
							    [bytes_array, path_array]							,					['''bytes''', '''path''']							,					mask=bytes_array.is_null()		)
				return array_cast(_A							,					self.pa_type		)







	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				@no_op_if_value_is_null
				def path_to_bytes(_A		):
							with xopen(_A							,					'''rb'''		) as f:
										UpperCAmelCase   		=				f.read()
							return bytes_

				UpperCAmelCase   		=				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()							,					)
				UpperCAmelCase   		=				pa.array(
				    [os.path.basename(_A		) if path is not None else None for path in storage.field('''path'''		).to_pylist()]							,					type=pa.string()							,					)
				UpperCAmelCase   		=				pa.StructArray.from_arrays([bytes_array, path_array]							,					['''bytes''', '''path''']							,					mask=bytes_array.is_null()		)
				return array_cast(_A							,					self.pa_type		)





def 						__SCREAMING_SNAKE_CASE  (				)  ->				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()
						UpperCAmelCase   		=				list(set(PIL.Image.OPEN.keys() ) & set(PIL.Image.SAVE.keys() ) )
			return _IMAGE_COMPRESSION_FORMATS





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				bytes:







			'''simple docstring'''

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





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				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 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				dict:







			'''simple docstring'''

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

			UpperCAmelCase   		=				array.dtype
			UpperCAmelCase   		=				dtype.byteorder if dtype.byteorder != '''=''' else _NATIVE_BYTEORDER
			UpperCAmelCase   		=				dtype.kind
			UpperCAmelCase   		=				dtype.itemsize

			UpperCAmelCase   		=				None

			# Multi-channel array case (only np.dtype("|u1") is allowed)
			if array.shape[2:]:
						UpperCAmelCase   		=				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:
						UpperCAmelCase   		=				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:
									UpperCAmelCase   		=				dtype_byteorder + dtype_kind + str(UpperCamelCase__ )
									UpperCAmelCase   		=				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}""" )
			UpperCAmelCase   		=				PIL.Image.fromarray(array.astype(UpperCamelCase__ ) )
			return {"path": None, "bytes": image_to_bytes(UpperCamelCase__ )}






def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				List[dict]:







			'''simple docstring'''

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

			if objs:
						UpperCAmelCase		,			UpperCAmelCase   		=				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 ):
									UpperCAmelCase   		=				no_op_if_value_is_null(UpperCamelCase__ )
									return [obj_to_image_dict_func(UpperCamelCase__ ) for obj in objs]
						elif isinstance(UpperCamelCase__			,			PIL.Image.Image ):
									UpperCAmelCase   		=				no_op_if_value_is_null(UpperCamelCase__ )
									return [obj_to_image_dict_func(UpperCamelCase__ ) for obj in objs]
						else:
									return objs
			else:
						return objs










import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np

class 	A_       :
	def      _lowercase				(    self		):






				'''simple docstring'''







				torch.manual_seed(0		)
				UpperCAmelCase   		=				TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				UNetaDConditionModel(
				    sample_size=3_2							,					layers_per_block=1							,					block_out_channels=[3_2, 6_4]							,					down_block_types=[
				        '''ResnetDownsampleBlock2D''',
				        '''SimpleCrossAttnDownBlock2D''',
				    ]							,					mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''							,					up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']							,					in_channels=3							,					out_channels=6							,					cross_attention_dim=3_2							,					encoder_hid_dim=3_2							,					attention_head_dim=8							,					addition_embed_type='''text'''							,					addition_embed_type_num_heads=2							,					cross_attention_norm='''group_norm'''							,					resnet_time_scale_shift='''scale_shift'''							,					act_fn='''gelu'''							,					)
				unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

				torch.manual_seed(0		)
				UpperCAmelCase   		=				DDPMScheduler(
				    num_train_timesteps=1_0_0_0							,					beta_schedule='''squaredcos_cap_v2'''							,					beta_start=0.00_01							,					beta_end=0.02							,					thresholding=_A							,					dynamic_thresholding_ratio=0.95							,					sample_max_value=1.0							,					prediction_type='''epsilon'''							,					variance_type='''learned_range'''							,					)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				IFWatermarker()

				return {
				    "text_encoder": text_encoder,
				    "tokenizer": tokenizer,
				    "unet": unet,
				    "scheduler": scheduler,
				    "watermarker": watermarker,
				    "safety_checker": None,
				    "feature_extractor": None,
				}
	def      _lowercase				(    self		):






				'''simple docstring'''







				torch.manual_seed(0		)
				UpperCAmelCase   		=				TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				UNetaDConditionModel(
				    sample_size=3_2							,					layers_per_block=[1, 2]							,					block_out_channels=[3_2, 6_4]							,					down_block_types=[
				        '''ResnetDownsampleBlock2D''',
				        '''SimpleCrossAttnDownBlock2D''',
				    ]							,					mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''							,					up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']							,					in_channels=6							,					out_channels=6							,					cross_attention_dim=3_2							,					encoder_hid_dim=3_2							,					attention_head_dim=8							,					addition_embed_type='''text'''							,					addition_embed_type_num_heads=2							,					cross_attention_norm='''group_norm'''							,					resnet_time_scale_shift='''scale_shift'''							,					act_fn='''gelu'''							,					class_embed_type='''timestep'''							,					mid_block_scale_factor=1.4_14							,					time_embedding_act_fn='''gelu'''							,					time_embedding_dim=3_2							,					)
				unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

				torch.manual_seed(0		)
				UpperCAmelCase   		=				DDPMScheduler(
				    num_train_timesteps=1_0_0_0							,					beta_schedule='''squaredcos_cap_v2'''							,					beta_start=0.00_01							,					beta_end=0.02							,					thresholding=_A							,					dynamic_thresholding_ratio=0.95							,					sample_max_value=1.0							,					prediction_type='''epsilon'''							,					variance_type='''learned_range'''							,					)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				DDPMScheduler(
				    num_train_timesteps=1_0_0_0							,					beta_schedule='''squaredcos_cap_v2'''							,					beta_start=0.00_01							,					beta_end=0.02							,					)

				torch.manual_seed(0		)
				UpperCAmelCase   		=				IFWatermarker()

				return {
				    "text_encoder": text_encoder,
				    "tokenizer": tokenizer,
				    "unet": unet,
				    "scheduler": scheduler,
				    "image_noising_scheduler": image_noising_scheduler,
				    "watermarker": watermarker,
				    "safety_checker": None,
				    "feature_extractor": None,
				}
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.get_dummy_components()
				UpperCAmelCase   		=				self.pipeline_class(**_A		)
				pipe.to(_A		)
				pipe.set_progress_bar_config(disable=_A		)

				UpperCAmelCase   		=				self.get_dummy_inputs(_A		)

				UpperCAmelCase   		=				inputs['''prompt''']
				UpperCAmelCase   		=				inputs['''generator''']
				UpperCAmelCase   		=				inputs['''num_inference_steps''']
				UpperCAmelCase   		=				inputs['''output_type''']

				if "image" in inputs:
							UpperCAmelCase   		=				inputs['''image''']
				else:
							UpperCAmelCase   		=				None

				if "mask_image" in inputs:
							UpperCAmelCase   		=				inputs['''mask_image''']
				else:
							UpperCAmelCase   		=				None

				if "original_image" in inputs:
							UpperCAmelCase   		=				inputs['''original_image''']
				else:
							UpperCAmelCase   		=				None

				UpperCAmelCase		,			UpperCAmelCase   		=				pipe.encode_prompt(_A		)

				# inputs with prompt converted to embeddings
				UpperCAmelCase   		=				{
				    '''prompt_embeds''': prompt_embeds,
				    '''negative_prompt_embeds''': negative_prompt_embeds,
				    '''generator''': generator,
				    '''num_inference_steps''': num_inference_steps,
				    '''output_type''': output_type,
				}

				if image is not None:
							UpperCAmelCase   		=				image

				if mask_image is not None:
							UpperCAmelCase   		=				mask_image

				if original_image is not None:
							UpperCAmelCase   		=				original_image

				# set all optional components to None
				for optional_component in pipe._optional_components:
							setattr(_A							,					_A							,					_A		)

				UpperCAmelCase   		=				pipe(**_A		)[0]

				with tempfile.TemporaryDirectory() as tmpdir:
							pipe.save_pretrained(_A		)
							UpperCAmelCase   		=				self.pipeline_class.from_pretrained(_A		)
							pipe_loaded.to(_A		)
							pipe_loaded.set_progress_bar_config(disable=_A		)

				pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

				for optional_component in pipe._optional_components:
							self.assertTrue(
							    getattr(_A							,					_A		) is None							,					F"""`{optional_component}` did not stay set to None after loading."""							,					)

				UpperCAmelCase   		=				self.get_dummy_inputs(_A		)

				UpperCAmelCase   		=				inputs['''generator''']
				UpperCAmelCase   		=				inputs['''num_inference_steps''']
				UpperCAmelCase   		=				inputs['''output_type''']

				# inputs with prompt converted to embeddings
				UpperCAmelCase   		=				{
				    '''prompt_embeds''': prompt_embeds,
				    '''negative_prompt_embeds''': negative_prompt_embeds,
				    '''generator''': generator,
				    '''num_inference_steps''': num_inference_steps,
				    '''output_type''': output_type,
				}

				if image is not None:
							UpperCAmelCase   		=				image

				if mask_image is not None:
							UpperCAmelCase   		=				mask_image

				if original_image is not None:
							UpperCAmelCase   		=				original_image

				UpperCAmelCase   		=				pipe_loaded(**_A		)[0]

				UpperCAmelCase   		=				np.abs(to_np(_A		) - to_np(_A		)		).max()
				self.assertLess(_A							,					1E-4		)







	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.get_dummy_components()
				UpperCAmelCase   		=				self.pipeline_class(**_A		)
				pipe.to(_A		)
				pipe.set_progress_bar_config(disable=_A		)

				UpperCAmelCase   		=				self.get_dummy_inputs(_A		)
				UpperCAmelCase   		=				pipe(**_A		)[0]

				with tempfile.TemporaryDirectory() as tmpdir:
							pipe.save_pretrained(_A		)
							UpperCAmelCase   		=				self.pipeline_class.from_pretrained(_A		)
							pipe_loaded.to(_A		)
							pipe_loaded.set_progress_bar_config(disable=_A		)

				pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()		)  # For reproducibility tests

				UpperCAmelCase   		=				self.get_dummy_inputs(_A		)
				UpperCAmelCase   		=				pipe_loaded(**_A		)[0]

				UpperCAmelCase   		=				np.abs(to_np(_A		) - to_np(_A		)		).max()
				self.assertLess(_A							,					1E-4		)



