






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class 							__UpperCamelCase     (  _lowerCAmelCase ,      _lowerCAmelCase ,      _lowerCAmelCase ,      unittest.TestCase ):
	__snake_case						:Optional[int] =							StableDiffusionInstructPixaPixPipeline
	__snake_case						:List[Any] =							TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width', 'cross_attention_kwargs'}
	__snake_case						:Any =							TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
	__snake_case						:int =							IMAGE_TO_IMAGE_IMAGE_PARAMS
	__snake_case						:int =							IMAGE_TO_IMAGE_IMAGE_PARAMS







	def 			_a (	self   :				List[Any]      )		->			Tuple:
		"""simple docstring"""







		torch.manual_seed(0      )
		__lowercase   							=	UNetaDConditionModel(
		    block_out_channels=(32, 64)							,					layers_per_block=2							,					sample_size=32							,					in_channels=8							,					out_channels=4							,					down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")							,					up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")							,					cross_attention_dim=32							,					)
		__lowercase   							=	PNDMScheduler(skip_prk_steps=_lowerCAmelCase      )
		torch.manual_seed(0      )
		__lowercase   							=	AutoencoderKL(
		    block_out_channels=[32, 64]							,					in_channels=3							,					out_channels=3							,					down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]							,					up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]							,					latent_channels=4							,					)
		torch.manual_seed(0      )
		__lowercase   							=	CLIPTextConfig(
		    bos_token_id=0							,					eos_token_id=2							,					hidden_size=32							,					intermediate_size=37							,					layer_norm_eps=1e-05							,					num_attention_heads=4							,					num_hidden_layers=5							,					pad_token_id=1							,					vocab_size=1000							,					)
		__lowercase   							=	CLIPTextModel(_lowerCAmelCase      )
		__lowercase   							=	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""      )

		__lowercase   							=	{
		    """unet""": unet,
		    """scheduler""": scheduler,
		    """vae""": vae,
		    """text_encoder""": text_encoder,
		    """tokenizer""": tokenizer,
		    """safety_checker""": None,
		    """feature_extractor""": None,
		}
		return components







	def 			_a (	self   :				Optional[int]							,					_lowerCAmelCase   :				str							,					_lowerCAmelCase   :				Optional[Any]=0      )		->			Any:
		"""simple docstring"""







		__lowercase   							=	floats_tensor((1, 3, 32, 32)							,					rng=random.Random(_lowerCAmelCase      )      ).to(_lowerCAmelCase      )
		__lowercase   							=	image.cpu().permute(0							,					2							,					3							,					1      )[0]
		__lowercase   							=	Image.fromarray(np.uinta(_lowerCAmelCase      )      ).convert("""RGB"""      )
		if str(_lowerCAmelCase      ).startswith("""mps"""      ):
			__lowercase   							=	torch.manual_seed(_lowerCAmelCase      )
		else:
			__lowercase   							=	torch.Generator(device=_lowerCAmelCase      ).manual_seed(_lowerCAmelCase      )
		__lowercase   							=	{
		    """prompt""": """A painting of a squirrel eating a burger""",
		    """image""": image,
		    """generator""": generator,
		    """num_inference_steps""": 2,
		    """guidance_scale""": 6.0,
		    """image_guidance_scale""": 1,
		    """output_type""": """numpy""",
		}
		return inputs







	def 			_a (	self   :				str      )		->			Tuple:
		"""simple docstring"""







		__lowercase   							=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
		__lowercase   							=	self.get_dummy_components()
		__lowercase   							=	StableDiffusionInstructPixaPixPipeline(**_lowerCAmelCase      )
		__lowercase   							=	sd_pipe.to(_lowerCAmelCase      )
		sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

		__lowercase   							=	self.get_dummy_inputs(_lowerCAmelCase      )
		__lowercase   							=	sd_pipe(**_lowerCAmelCase      ).images
		__lowercase   							=	image[0, -3:, -3:, -1]
		assert image.shape == (1, 32, 32, 3)
		__lowercase   							=	np.array([0.7_526, 0.3_750, 0.4_547, 0.6_117, 0.5_866, 0.5_016, 0.4_327, 0.5_642, 0.4_815]      )

		assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-3







	def 			_a (	self   :				Any      )		->			Any:
		"""simple docstring"""







		__lowercase   							=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
		__lowercase   							=	self.get_dummy_components()
		__lowercase   							=	StableDiffusionInstructPixaPixPipeline(**_lowerCAmelCase      )
		__lowercase   							=	sd_pipe.to(_lowerCAmelCase      )
		sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

		__lowercase   							=	self.get_dummy_inputs(_lowerCAmelCase      )
		__lowercase   							=	"""french fries"""
		__lowercase   							=	sd_pipe(**_lowerCAmelCase							,					negative_prompt=_lowerCAmelCase      )
		__lowercase   							=	output.images
		__lowercase   							=	image[0, -3:, -3:, -1]

		assert image.shape == (1, 32, 32, 3)
		__lowercase   							=	np.array([0.7_511, 0.3_642, 0.4_553, 0.6_236, 0.5_797, 0.5_013, 0.4_343, 0.5_611, 0.4_831]      )

		assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-3







	def 			_a (	self   :				Optional[Any]      )		->			Dict:
		"""simple docstring"""







		__lowercase   							=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
		__lowercase   							=	self.get_dummy_components()
		__lowercase   							=	StableDiffusionInstructPixaPixPipeline(**_lowerCAmelCase      )
		__lowercase   							=	sd_pipe.to(_lowerCAmelCase      )
		sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

		__lowercase   							=	self.get_dummy_inputs(_lowerCAmelCase      )
		__lowercase   							=	[inputs["""prompt"""]] * 2

		__lowercase   							=	np.array(inputs["""image"""]      ).astype(np.floataa      ) / 255.0
		__lowercase   							=	torch.from_numpy(_lowerCAmelCase      ).unsqueeze(0      ).to(_lowerCAmelCase      )
		__lowercase   							=	image / 2 + 0.5
		__lowercase   							=	image.permute(0							,					3							,					1							,					2      )
		__lowercase   							=	image.repeat(2							,					1							,					1							,					1      )

		__lowercase   							=	sd_pipe(**_lowerCAmelCase      ).images
		__lowercase   							=	image[-1, -3:, -3:, -1]

		assert image.shape == (2, 32, 32, 3)
		__lowercase   							=	np.array([0.5_812, 0.5_748, 0.5_222, 0.5_908, 0.5_695, 0.7_174, 0.6_804, 0.5_523, 0.5_579]      )

		assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-3







	def 			_a (	self   :				int      )		->			List[str]:
		"""simple docstring"""







		__lowercase   							=	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
		__lowercase   							=	self.get_dummy_components()
		__lowercase   							=	EulerAncestralDiscreteScheduler(
		    beta_start=0.00_085							,					beta_end=0.012							,					beta_schedule="""scaled_linear"""      )
		__lowercase   							=	StableDiffusionInstructPixaPixPipeline(**_lowerCAmelCase      )
		__lowercase   							=	sd_pipe.to(_lowerCAmelCase      )
		sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

		__lowercase   							=	self.get_dummy_inputs(_lowerCAmelCase      )
		__lowercase   							=	sd_pipe(**_lowerCAmelCase      ).images
		__lowercase   							=	image[0, -3:, -3:, -1]

		__lowercase   							=	[round(_lowerCAmelCase							,					4      ) for x in image_slice.flatten().tolist()]
		print(""",""".join([str(_lowerCAmelCase      ) for x in slice]      )      )

		assert image.shape == (1, 32, 32, 3)
		__lowercase   							=	np.array([0.7_417, 0.3_842, 0.4_732, 0.5_776, 0.5_891, 0.5_139, 0.4_052, 0.5_673, 0.4_986]      )

		assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-3







	def 			_a (	self   :				List[str]      )		->			Dict:
		"""simple docstring"""







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







	def 			_a (	self   :				Union[str, Any]      )		->			int:
		"""simple docstring"""







		__lowercase   							=	self.get_dummy_components()
		__lowercase   							=	StableDiffusionInstructPixaPixPipeline(**_lowerCAmelCase      )
		__lowercase   							=	VaeImageProcessor(do_resize=_lowerCAmelCase							,					do_normalize=_lowerCAmelCase      )
		__lowercase   							=	pipe.to(_lowerCAmelCase      )
		pipe.set_progress_bar_config(disable=_lowerCAmelCase      )

		__lowercase   							=	pipe(**self.get_dummy_inputs_by_type(_lowerCAmelCase							,					input_image_type="""pt"""      )      )[0]

		__lowercase   							=	components["""vae"""]
		__lowercase   							=	self.get_dummy_inputs_by_type(_lowerCAmelCase							,					input_image_type="""pt"""      )

		for image_param in self.image_latents_params:
			if image_param in inputs.keys():
				__lowercase   							=	vae.encode(inputs[image_param]      ).latent_dist.mode()

		__lowercase   							=	pipe(**_lowerCAmelCase      )[0]

		__lowercase   							=	np.abs(out - out_latents_inputs      ).max()
		self.assertLess(_lowerCAmelCase							,					1e-4							,					"""passing latents as image input generate different result from passing image"""      )


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







	def 			_a (	self   :				Dict      )		->			Optional[int]:
		"""simple docstring"""







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







	def 			_a (	self   :				Dict							,					_lowerCAmelCase   :				str=0      )		->			Optional[int]:
		"""simple docstring"""







		__lowercase   							=	torch.manual_seed(_lowerCAmelCase      )
		__lowercase   							=	load_image(
		    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg"""      )
		__lowercase   							=	{
		    """prompt""": """turn him into a cyborg""",
		    """image""": image,
		    """generator""": generator,
		    """num_inference_steps""": 3,
		    """guidance_scale""": 7.5,
		    """image_guidance_scale""": 1.0,
		    """output_type""": """numpy""",
		}
		return inputs







	def 			_a (	self   :				Tuple      )		->			Optional[Any]:
		"""simple docstring"""







		__lowercase   							=	StableDiffusionInstructPixaPixPipeline.from_pretrained(
		    """timbrooks/instruct-pix2pix"""							,					safety_checker=_lowerCAmelCase      )
		pipe.to(_lowerCAmelCase      )
		pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
		pipe.enable_attention_slicing()

		__lowercase   							=	self.get_inputs()
		__lowercase   							=	pipe(**_lowerCAmelCase      ).images
		__lowercase   							=	image[0, -3:, -3:, -1].flatten()

		assert image.shape == (1, 512, 512, 3)
		__lowercase   							=	np.array([0.5_902, 0.6_015, 0.6_027, 0.5_983, 0.6_092, 0.6_061, 0.5_765, 0.5_785, 0.5_555]      )

		assert np.abs(expected_slice - image_slice      ).max() < 1e-3







	def 			_a (	self   :				Dict      )		->			Union[str, Any]:
		"""simple docstring"""







		__lowercase   							=	StableDiffusionInstructPixaPixPipeline.from_pretrained(
		    """timbrooks/instruct-pix2pix"""							,					safety_checker=_lowerCAmelCase      )
		__lowercase   							=	LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
		pipe.to(_lowerCAmelCase      )
		pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
		pipe.enable_attention_slicing()

		__lowercase   							=	self.get_inputs()
		__lowercase   							=	pipe(**_lowerCAmelCase      ).images
		__lowercase   							=	image[0, -3:, -3:, -1].flatten()

		assert image.shape == (1, 512, 512, 3)
		__lowercase   							=	np.array([0.6_578, 0.6_817, 0.6_972, 0.6_761, 0.6_856, 0.6_916, 0.6_428, 0.6_516, 0.6_301]      )

		assert np.abs(expected_slice - image_slice      ).max() < 1e-3







	def 			_a (	self   :				Any      )		->			Union[str, Any]:
		"""simple docstring"""







		__lowercase   							=	StableDiffusionInstructPixaPixPipeline.from_pretrained(
		    """timbrooks/instruct-pix2pix"""							,					safety_checker=_lowerCAmelCase      )
		__lowercase   							=	DDIMScheduler.from_config(pipe.scheduler.config      )
		pipe.to(_lowerCAmelCase      )
		pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
		pipe.enable_attention_slicing()

		__lowercase   							=	self.get_inputs()
		__lowercase   							=	pipe(**_lowerCAmelCase      ).images
		__lowercase   							=	image[0, -3:, -3:, -1].flatten()

		assert image.shape == (1, 512, 512, 3)
		__lowercase   							=	np.array([0.3_828, 0.3_834, 0.3_818, 0.3_792, 0.3_865, 0.3_752, 0.3_792, 0.3_847, 0.3_753]      )

		assert np.abs(expected_slice - image_slice      ).max() < 1e-3







	def 			_a (	self   :				Union[str, Any]      )		->			Any:
		"""simple docstring"""







		__lowercase   							=	0

		def callback_fn(_lowerCAmelCase   :				int							,					_lowerCAmelCase   :				int							,					_lowerCAmelCase   :				torch.FloatTensor      ) -> None:
			__lowercase   							=	True
			nonlocal number_of_steps
			number_of_steps += 1
			if step == 1:
				__lowercase   							=	latents.detach().cpu().numpy()
				assert latents.shape == (1, 4, 64, 64)
				__lowercase   							=	latents[0, -3:, -3:, -1]
				__lowercase   							=	np.array([-0.2_463, -0.4_644, -0.9_756, 1.5_176, 1.4_414, 0.7_866, 0.9_897, 0.8_521, 0.7_983]      )

				assert np.abs(latents_slice.flatten() - expected_slice      ).max() < 5e-2
			elif step == 2:
				__lowercase   							=	latents.detach().cpu().numpy()
				assert latents.shape == (1, 4, 64, 64)
				__lowercase   							=	latents[0, -3:, -3:, -1]
				__lowercase   							=	np.array([-0.2_644, -0.4_626, -0.9_653, 1.5_176, 1.4_551, 0.7_686, 0.9_805, 0.8_452, 0.8_115]      )

				assert np.abs(latents_slice.flatten() - expected_slice      ).max() < 5e-2

		__lowercase   							=	False

		__lowercase   							=	StableDiffusionInstructPixaPixPipeline.from_pretrained(
		    """timbrooks/instruct-pix2pix"""							,					safety_checker=_lowerCAmelCase							,					torch_dtype=torch.floataa      )
		__lowercase   							=	pipe.to(_lowerCAmelCase      )
		pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
		pipe.enable_attention_slicing()

		__lowercase   							=	self.get_inputs()
		pipe(**_lowerCAmelCase							,					callback=_lowerCAmelCase							,					callback_steps=1      )
		assert callback_fn.has_been_called
		assert number_of_steps == 3







	def 			_a (	self   :				Optional[int]      )		->			str:
		"""simple docstring"""







		torch.cuda.empty_cache()
		torch.cuda.reset_max_memory_allocated()
		torch.cuda.reset_peak_memory_stats()

		__lowercase   							=	StableDiffusionInstructPixaPixPipeline.from_pretrained(
		    """timbrooks/instruct-pix2pix"""							,					safety_checker=_lowerCAmelCase							,					torch_dtype=torch.floataa      )
		__lowercase   							=	pipe.to(_lowerCAmelCase      )
		pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
		pipe.enable_attention_slicing(1      )
		pipe.enable_sequential_cpu_offload()

		__lowercase   							=	self.get_inputs()
		__lowercase   							=	pipe(**_lowerCAmelCase      )

		__lowercase   							=	torch.cuda.max_memory_allocated()
		# make sure that less than 2.2 GB is allocated
		assert mem_bytes < 2.2 * 10**9







	def 			_a (	self   :				List[Any]      )		->			Optional[int]:
		"""simple docstring"""







		__lowercase   							=	self.get_inputs()
		# resize to resolution that is divisible by 8 but not 16 or 32
		__lowercase   							=	inputs["""image"""].resize((504, 504)      )

		__lowercase   							=	"""timbrooks/instruct-pix2pix"""
		__lowercase   							=	StableDiffusionInstructPixaPixPipeline.from_pretrained(
		    _lowerCAmelCase							,					safety_checker=_lowerCAmelCase							,					)
		pipe.to(_lowerCAmelCase      )
		pipe.set_progress_bar_config(disable=_lowerCAmelCase      )
		pipe.enable_attention_slicing()

		__lowercase   							=	pipe(**_lowerCAmelCase      )
		__lowercase   							=	output.images[0]

		__lowercase   							=	image[255:258, 383:386, -1]

		assert image.shape == (504, 504, 3)
		__lowercase   							=	np.array([0.2_726, 0.2_529, 0.2_664, 0.2_655, 0.2_641, 0.2_642, 0.2_591, 0.2_649, 0.2_590]      )

		assert np.abs(image_slice.flatten() - expected_slice      ).max() < 5e-3



'''simple docstring'''



from sklearn.metrics import mean_squared_error

import datasets


_lowercase          =						"""\
@article{scikit-learn,
 title={Scikit-learn: Machine Learning in {P}ython},
 author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.
         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.
         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and
         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
 journal={Journal of Machine Learning Research},
 volume={12},
 pages={2825--2830},
 year={2011}
}
"""

_lowercase          =						"""\
Mean Squared Error(MSE) is the average of the square of difference between the predicted
and actual values.
"""


_lowercase          =						"""
Args:
    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Estimated target values.
    references: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Ground truth (correct) target values.
    sample_weight: array-like of shape (n_samples,), default=None
        Sample weights.
    multioutput: {\"raw_values\", \"uniform_average\"} or array-like of shape (n_outputs,), default=\"uniform_average\"
        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.

        \"raw_values\" : Returns a full set of errors in case of multioutput input.

        \"uniform_average\" : Errors of all outputs are averaged with uniform weight.

    squared : bool, default=True
        If True returns MSE value, if False returns RMSE (Root Mean Squared Error) value.

Returns:
    mse : mean squared error.
Examples:

    >>> mse_metric = datasets.load_metric(\"mse\")
    >>> predictions = [2.5, 0.0, 2, 8]
    >>> references = [3, -0.5, 2, 7]
    >>> results = mse_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'mse': 0.375}
    >>> rmse_result = mse_metric.compute(predictions=predictions, references=references, squared=False)
    >>> print(rmse_result)
    {'mse': 0.6123724356957945}

    If you're using multi-dimensional lists, then set the config as follows :

    >>> mse_metric = datasets.load_metric(\"mse\", \"multilist\")
    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]
    >>> references = [[0, 2], [-1, 2], [8, -5]]
    >>> results = mse_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'mse': 0.7083333333333334}
    >>> results = mse_metric.compute(predictions=predictions, references=references, multioutput='raw_values')
    >>> print(results) # doctest: +NORMALIZE_WHITESPACE
    {'mse': array([0.41666667, 1. ])}
"""




@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION			,				_KWARGS_DESCRIPTION					)
class     UpperCAmelCase_	(						datasets.Metric					):






						'''simple docstring'''






						def    _lowercase      (   self ):

												"""simple docstring"""



												return datasets.MetricInfo(
												    description=_DESCRIPTION       ,  citation=_CITATION       ,  inputs_description=_KWARGS_DESCRIPTION       ,  features=datasets.Features(self._get_feature_types() )       ,  reference_urls=[
												        """https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html"""
												    ]       ,  )




						def    _lowercase      (   self ):

												"""simple docstring"""



												if self.config_name == "multilist":
																		return {
																		    "predictions": datasets.Sequence(datasets.Value("""float""" ) ),
																		    "references": datasets.Sequence(datasets.Value("""float""" ) ),
																		}
												else:
																		return {
																		    "predictions": datasets.Value("""float""" ),
																		    "references": datasets.Value("""float""" ),
																		}




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase       ,  _lowercase=None       ,  _lowercase="uniform_average"       ,  _lowercase=True ):

												"""simple docstring"""



												_lowerCAmelCase    		=					mean_squared_error(
												    _lowercase       ,  _lowercase       ,  sample_weight=_lowercase       ,  multioutput=_lowercase       ,  squared=_lowercase )

												return {"mse": mse}
