



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, is_vision_available, logging


if is_vision_available():
	import PIL


snake_case__					:     Optional[Any]    		=   logging.get_logger(__name__)

class   _A   ( _UpperCamelCase	):


						'''simple docstring'''





						_snake_case			:     Optional[int]      							=      ["""pixel_values"""]



						def __init__(	self	:							str							,      lowerCamelCase	:							Tuple = True							,      lowerCamelCase	:							List[str] = None							,      lowerCamelCase	:							Tuple = None							,      lowerCamelCase	:							Optional[Any] = PILImageResampling.BILINEAR							,      lowerCamelCase	:							Union[str, Any] = True							,      lowerCamelCase	:							str = 1 / 255							,      lowerCamelCase	:							Optional[Any] = True							,      lowerCamelCase	:							int = None							,      lowerCamelCase	:							List[str] = None							,      **lowerCamelCase	:							str							,      ):






							'''simple docstring'''



							super().__init__(**_UpperCAmelCase	)
							__lowercase             =			size if size is not None else {'''shortest_edge''': 384}
							__lowercase             =			get_size_dict(_UpperCAmelCase							,      default_to_square=_UpperCAmelCase	)

							__lowercase             =			do_resize
							__lowercase             =			size
							# Default value set here for backwards compatibility where the value in config is None
							__lowercase             =			crop_pct if crop_pct is not None else 224 / 256
							__lowercase             =			resample
							__lowercase             =			do_rescale
							__lowercase             =			rescale_factor
							__lowercase             =			do_normalize
							__lowercase             =			image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
							__lowercase             =			image_std if image_std is not None else IMAGENET_STANDARD_STD



						def     _snake_case	(	self	:							int							,      lowerCamelCase	:							str							,      lowerCamelCase	:							Optional[Any]							,      lowerCamelCase	:							Any							,      lowerCamelCase	:							Optional[int] = PILImageResampling.BICUBIC							,      lowerCamelCase	:							Union[str, Any] = None							,      **lowerCamelCase	:							Any							,      ):






							'''simple docstring'''



							__lowercase             =			get_size_dict(_UpperCAmelCase							,      default_to_square=_UpperCAmelCase	)
							if "shortest_edge" not in size:
								raise ValueError(f"""Size dictionary must contain 'shortest_edge' key. Got {size.keys()}"""	)
							__lowercase             =			size['''shortest_edge''']

							if shortest_edge < 384:
								# maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
								__lowercase             =			int(shortest_edge / crop_pct	)
								__lowercase             =			get_resize_output_image_size(_UpperCAmelCase							,      size=_UpperCAmelCase							,      default_to_square=_UpperCAmelCase	)
								__lowercase             =			resize(image=_UpperCAmelCase							,      size=_UpperCAmelCase							,      resample=_UpperCAmelCase							,      data_format=_UpperCAmelCase							,      **_UpperCAmelCase	)
								# then crop to (shortest_edge, shortest_edge)
								return center_crop(image=_UpperCAmelCase							,      size=(shortest_edge, shortest_edge)							,      data_format=_UpperCAmelCase							,      **_UpperCAmelCase	)
							else:
								# warping (no cropping) when evaluated at 384 or larger
								return resize(
								    _UpperCAmelCase							,      size=(shortest_edge, shortest_edge)							,      resample=_UpperCAmelCase							,      data_format=_UpperCAmelCase							,      **_UpperCAmelCase	)



						def     _snake_case	(	self	:							Tuple							,      lowerCamelCase	:							Union[str, Any]							,      lowerCamelCase	:							Union[str, Any]							,      lowerCamelCase	:							Optional[int] = None							,      **lowerCamelCase	:							int							,      ):






							'''simple docstring'''



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



						def     _snake_case	(	self	:							Dict							,      lowerCamelCase	:							Dict							,      lowerCamelCase	:							Union[str, Any]							,      lowerCamelCase	:							int							,      lowerCamelCase	:							Dict = None							,      **lowerCamelCase	:							List[str]							,      ):






							'''simple docstring'''



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



						def     _snake_case	(	self	:							int							,      lowerCamelCase	:							Tuple							,      lowerCamelCase	:							Optional[int] = None							,      lowerCamelCase	:							Tuple = None							,      lowerCamelCase	:							Optional[int] = None							,      lowerCamelCase	:							str = None							,      lowerCamelCase	:							Any = None							,      lowerCamelCase	:							Optional[Any] = None							,      lowerCamelCase	:							Dict = None							,      lowerCamelCase	:							Optional[Any] = None							,      lowerCamelCase	:							List[str] = None							,      lowerCamelCase	:							Optional[int] = None							,      lowerCamelCase	:							Union[str, Any] = ChannelDimension.FIRST							,      **lowerCamelCase	:							Optional[Any]							,      ):






							'''simple docstring'''



							__lowercase             =			do_resize if do_resize is not None else self.do_resize
							__lowercase             =			crop_pct if crop_pct is not None else self.crop_pct
							__lowercase             =			resample if resample is not None else self.resample
							__lowercase             =			do_rescale if do_rescale is not None else self.do_rescale
							__lowercase             =			rescale_factor if rescale_factor is not None else self.rescale_factor
							__lowercase             =			do_normalize if do_normalize is not None else self.do_normalize
							__lowercase             =			image_mean if image_mean is not None else self.image_mean
							__lowercase             =			image_std if image_std is not None else self.image_std

							__lowercase             =			size if size is not None else self.size
							__lowercase             =			get_size_dict(_UpperCAmelCase							,      default_to_square=_UpperCAmelCase	)

							__lowercase             =			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 or resample is None:
								raise ValueError("Size and resample must be specified if do_resize is True."	)

							if do_resize and size["shortest_edge"] < 384 and crop_pct is None:
								raise ValueError("crop_pct must be specified if size < 384."	)

							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.
							__lowercase             =			[to_numpy_array(_UpperCAmelCase	) for image in images]

							if do_resize:
								__lowercase             =			[self.resize(image=_UpperCAmelCase							,      size=_UpperCAmelCase							,      crop_pct=_UpperCAmelCase							,      resample=_UpperCAmelCase	) for image in images]

							if do_rescale:
								__lowercase             =			[self.rescale(image=_UpperCAmelCase							,      scale=_UpperCAmelCase	) for image in images]

							if do_normalize:
								__lowercase             =			[self.normalize(image=_UpperCAmelCase							,      mean=_UpperCAmelCase							,      std=_UpperCAmelCase	) for image in images]

							__lowercase             =			[to_channel_dimension_format(_UpperCAmelCase							,      _UpperCAmelCase	) for image in images]

							__lowercase             =			{'''pixel_values''': images}
							return BatchFeature(data=_UpperCAmelCase							,      tensor_type=_UpperCAmelCase	)








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   :


						'''simple docstring'''





						def     _snake_case	(	self	:							Union[str, Any]	):






							'''simple docstring'''



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

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

							torch.manual_seed(0	)
							__lowercase		       =			UNetaDConditionModel(
							    sample_size=32							,      layers_per_block=1							,      block_out_channels=[32, 64]							,      down_block_types=[
							        "ResnetDownsampleBlock2D",
							        "SimpleCrossAttnDownBlock2D",
							    ]							,      mid_block_type="UNetMidBlock2DSimpleCrossAttn"							,      up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]							,      in_channels=3							,      out_channels=6							,      cross_attention_dim=32							,      encoder_hid_dim=32							,      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	)
							__lowercase		       =			DDPMScheduler(
							    num_train_timesteps=1_000							,      beta_schedule="squaredcos_cap_v2"							,      beta_start=0.0001							,      beta_end=0.02							,      thresholding=lowerCamelCase							,      dynamic_thresholding_ratio=0.95							,      sample_max_value=1.0							,      prediction_type="epsilon"							,      variance_type="learned_range"							,      )

							torch.manual_seed(0	)
							__lowercase		       =			IFWatermarker()

							return {
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "unet": unet,
							    "scheduler": scheduler,
							    "watermarker": watermarker,
							    "safety_checker": None,
							    "feature_extractor": None,
							}



						def     _snake_case	(	self	:							Tuple	):






							'''simple docstring'''



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

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

							torch.manual_seed(0	)
							__lowercase		       =			UNetaDConditionModel(
							    sample_size=32							,      layers_per_block=[1, 2]							,      block_out_channels=[32, 64]							,      down_block_types=[
							        "ResnetDownsampleBlock2D",
							        "SimpleCrossAttnDownBlock2D",
							    ]							,      mid_block_type="UNetMidBlock2DSimpleCrossAttn"							,      up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]							,      in_channels=6							,      out_channels=6							,      cross_attention_dim=32							,      encoder_hid_dim=32							,      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.414							,      time_embedding_act_fn="gelu"							,      time_embedding_dim=32							,      )
							unet.set_attn_processor(AttnAddedKVProcessor()	)  # For reproducibility tests

							torch.manual_seed(0	)
							__lowercase		       =			DDPMScheduler(
							    num_train_timesteps=1_000							,      beta_schedule="squaredcos_cap_v2"							,      beta_start=0.0001							,      beta_end=0.02							,      thresholding=lowerCamelCase							,      dynamic_thresholding_ratio=0.95							,      sample_max_value=1.0							,      prediction_type="epsilon"							,      variance_type="learned_range"							,      )

							torch.manual_seed(0	)
							__lowercase		       =			DDPMScheduler(
							    num_train_timesteps=1_000							,      beta_schedule="squaredcos_cap_v2"							,      beta_start=0.0001							,      beta_end=0.02							,      )

							torch.manual_seed(0	)
							__lowercase		       =			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     _snake_case	(	self	:							str	):






							'''simple docstring'''



							__lowercase		       =			self.get_dummy_components()
							__lowercase		       =			self.pipeline_class(**lowerCamelCase	)
							pipe.to(lowerCamelCase	)
							pipe.set_progress_bar_config(disable=lowerCamelCase	)

							__lowercase		       =			self.get_dummy_inputs(lowerCamelCase	)

							__lowercase		       =			inputs["prompt"]
							__lowercase		       =			inputs["generator"]
							__lowercase		       =			inputs["num_inference_steps"]
							__lowercase		       =			inputs["output_type"]

							if "image" in inputs:
								__lowercase		       =			inputs["image"]
							else:
								__lowercase		       =			None

							if "mask_image" in inputs:
								__lowercase		       =			inputs["mask_image"]
							else:
								__lowercase		       =			None

							if "original_image" in inputs:
								__lowercase		       =			inputs["original_image"]
							else:
								__lowercase		       =			None

							__lowercase       ,    __lowercase		       =			pipe.encode_prompt(lowerCamelCase	)

							# inputs with prompt converted to embeddings
							__lowercase		       =			{
							    "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:
								__lowercase		       =			image

							if mask_image is not None:
								__lowercase		       =			mask_image

							if original_image is not None:
								__lowercase		       =			original_image

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

							__lowercase		       =			pipe(**lowerCamelCase	)[0]

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

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

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

							__lowercase		       =			self.get_dummy_inputs(lowerCamelCase	)

							__lowercase		       =			inputs["generator"]
							__lowercase		       =			inputs["num_inference_steps"]
							__lowercase		       =			inputs["output_type"]

							# inputs with prompt converted to embeddings
							__lowercase		       =			{
							    "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:
								__lowercase		       =			image

							if mask_image is not None:
								__lowercase		       =			mask_image

							if original_image is not None:
								__lowercase		       =			original_image

							__lowercase		       =			pipe_loaded(**lowerCamelCase	)[0]

							__lowercase		       =			np.abs(to_np(lowerCamelCase	) - to_np(lowerCamelCase	)	).max()
							self.assertLess(lowerCamelCase							,      1e-4	)



						def     _snake_case	(	self	:							Optional[Any]	):






							'''simple docstring'''



							__lowercase		       =			self.get_dummy_components()
							__lowercase		       =			self.pipeline_class(**lowerCamelCase	)
							pipe.to(lowerCamelCase	)
							pipe.set_progress_bar_config(disable=lowerCamelCase	)

							__lowercase		       =			self.get_dummy_inputs(lowerCamelCase	)
							__lowercase		       =			pipe(**lowerCamelCase	)[0]

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

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

							__lowercase		       =			self.get_dummy_inputs(lowerCamelCase	)
							__lowercase		       =			pipe_loaded(**lowerCamelCase	)[0]

							__lowercase		       =			np.abs(to_np(lowerCamelCase	) - to_np(lowerCamelCase	)	).max()
							self.assertLess(lowerCamelCase							,      1e-4	)



