
'''simple docstring'''



import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
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     UpperCAmelCase_	(						_SCREAMING_SNAKE_CASE			,				unittest.TestCase					):






						'''simple docstring'''






						_lowercase						: List[Any]										=  KandinskyVaaImgaImgPipeline
						_lowercase						: List[Any]										=  ['''image_embeds''', '''negative_image_embeds''', '''image''']
						_lowercase						: str										=  [
						    '''image_embeds''',
						    '''negative_image_embeds''',
						    '''image''',
						]
						_lowercase						: int										=  [
						    '''generator''',
						    '''height''',
						    '''width''',
						    '''strength''',
						    '''guidance_scale''',
						    '''num_inference_steps''',
						    '''return_dict''',
						    '''guidance_scale''',
						    '''num_images_per_prompt''',
						    '''output_type''',
						    '''return_dict''',
						]
						_lowercase						: str										=  False




						@property
						def    _lowercase      (   self ):

												"""simple docstring"""



												return 32




						@property
						def    _lowercase      (   self ):

												"""simple docstring"""



												return 32




						@property
						def    _lowercase      (   self ):

												"""simple docstring"""



												return self.time_input_dim




						@property
						def    _lowercase      (   self ):

												"""simple docstring"""



												return self.time_input_dim * 4




						@property
						def    _lowercase      (   self ):

												"""simple docstring"""



												return 100




						@property
						def    _lowercase      (   self ):

												"""simple docstring"""



												torch.manual_seed(0 )

												_lowerCAmelCase    		=					{
												    """in_channels""": 4,
												    # Out channels is double in channels because predicts mean and variance
												    """out_channels""": 8,
												    """addition_embed_type""": """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""": """image_proj""",
												    """cross_attention_dim""": self.cross_attention_dim,
												    """attention_head_dim""": 4,
												    """resnet_time_scale_shift""": """scale_shift""",
												    """class_embed_type""": None,
												}

												_lowerCAmelCase    		=					UNetaDConditionModel(**_lowercase )
												return model




						@property
						def    _lowercase      (   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    _lowercase      (   self ):

												"""simple docstring"""



												torch.manual_seed(0 )
												_lowerCAmelCase    		=					VQModel(**self.dummy_movq_kwargs )
												return model




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					self.dummy_unet
												_lowerCAmelCase    		=					self.dummy_movq

												_lowerCAmelCase    		=					{
												    """num_train_timesteps""": 1_000,
												    """beta_schedule""": """linear""",
												    """beta_start""": 0.0_0085,
												    """beta_end""": 0.012,
												    """clip_sample""": False,
												    """set_alpha_to_one""": False,
												    """steps_offset""": 0,
												    """prediction_type""": """epsilon""",
												    """thresholding""": False,
												}

												_lowerCAmelCase    		=					DDIMScheduler(**_lowercase )

												_lowerCAmelCase    		=					{
												    """unet""": unet,
												    """scheduler""": scheduler,
												    """movq""": movq,
												}

												return components




						def    _lowercase      (   self       ,  _lowercase       ,  _lowercase=0 ):

												"""simple docstring"""



												_lowerCAmelCase    		=					floats_tensor((1, self.text_embedder_hidden_size)       ,  rng=random.Random(_lowercase ) ).to(_lowercase )
												_lowerCAmelCase    		=					floats_tensor((1, self.text_embedder_hidden_size)       ,  rng=random.Random(seed + 1 ) ).to(
												    _lowercase )
												# create init_image
												_lowerCAmelCase    		=					floats_tensor((1, 3, 64, 64)       ,  rng=random.Random(_lowercase ) ).to(_lowercase )
												_lowerCAmelCase    		=					image.cpu().permute(0       ,  2       ,  3       ,  1 )[0]
												_lowerCAmelCase    		=					Image.fromarray(np.uinta(_lowercase ) ).convert("""RGB""" ).resize((256, 256) )

												if str(_lowercase ).startswith("""mps""" ):
																		_lowerCAmelCase    		=					torch.manual_seed(_lowercase )
												else:
																		_lowerCAmelCase    		=					torch.Generator(device=_lowercase ).manual_seed(_lowercase )
												_lowerCAmelCase    		=					{
												    """image""": init_image,
												    """image_embeds""": image_embeds,
												    """negative_image_embeds""": negative_image_embeds,
												    """generator""": generator,
												    """height""": 64,
												    """width""": 64,
												    """num_inference_steps""": 10,
												    """guidance_scale""": 7.0,
												    """strength""": 0.2,
												    """output_type""": """np""",
												}
												return inputs




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					"""cpu"""

												_lowerCAmelCase    		=					self.get_dummy_components()

												_lowerCAmelCase    		=					self.pipeline_class(**_lowercase )
												_lowerCAmelCase    		=					pipe.to(_lowercase )

												pipe.set_progress_bar_config(disable=_lowercase )

												_lowerCAmelCase    		=					pipe(**self.get_dummy_inputs(_lowercase ) )
												_lowerCAmelCase    		=					output.images

												_lowerCAmelCase    		=					pipe(
												    **self.get_dummy_inputs(_lowercase )       ,  return_dict=_lowercase       ,  )[0]

												_lowerCAmelCase    		=					image[0, -3:, -3:, -1]
												_lowerCAmelCase    		=					image_from_tuple[0, -3:, -3:, -1]

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

												_lowerCAmelCase    		=					np.array(
												    [0.619_9778, 0.6398_4406, 0.4614_5785, 0.6294_4984, 0.562_2215, 0.4730_6132, 0.4744_1456, 0.460_7606, 0.4871_9263] )
												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()}'






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






						'''simple docstring'''






						def    _lowercase      (   self ):

												"""simple docstring"""



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




						def    _lowercase      (   self ):

												"""simple docstring"""



												_lowerCAmelCase    		=					load_numpy(
												    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
												    """/kandinskyv22/kandinskyv22_img2img_frog.npy""" )

												_lowerCAmelCase    		=					load_image(
												    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/kandinsky/cat.png""" )
												_lowerCAmelCase    		=					"""A red cartoon frog, 4k"""

												_lowerCAmelCase    		=					KandinskyVaaPriorPipeline.from_pretrained(
												    """kandinsky-community/kandinsky-2-2-prior"""       ,  torch_dtype=torch.floataa )
												pipe_prior.to(_lowercase )

												_lowerCAmelCase    		=					KandinskyVaaImgaImgPipeline.from_pretrained(
												    """kandinsky-community/kandinsky-2-2-decoder"""       ,  torch_dtype=torch.floataa )
												_lowerCAmelCase    		=					pipeline.to(_lowercase )

												pipeline.set_progress_bar_config(disable=_lowercase )

												_lowerCAmelCase    		=					torch.Generator(device="""cpu""" ).manual_seed(0 )
												_lowerCAmelCase		,       _lowerCAmelCase    		=					pipe_prior(
												    _lowercase       ,  generator=_lowercase       ,  num_inference_steps=5       ,  negative_prompt=""""""       ,  ).to_tuple()

												_lowerCAmelCase    		=					pipeline(
												    image=_lowercase       ,  image_embeds=_lowercase       ,  negative_image_embeds=_lowercase       ,  generator=_lowercase       ,  num_inference_steps=100       ,  height=768       ,  width=768       ,  strength=0.2       ,  output_type="""np"""       ,  )

												_lowerCAmelCase    		=					output.images[0]

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

												assert_mean_pixel_difference(_lowercase       ,  _lowercase )





'''simple docstring'''





from sklearn.metrics import mean_squared_error

import datasets


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

SCREAMING_SNAKE_CASE__          =  "\\nMean Squared Error(MSE) is the average of the square of difference between the predicted\nand actual values.\n"


SCREAMING_SNAKE_CASE__          =  "\nArgs:\n    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)\n        Estimated target values.\n    references: array-like of shape (n_samples,) or (n_samples, n_outputs)\n        Ground truth (correct) target values.\n    sample_weight: array-like of shape (n_samples,), default=None\n        Sample weights.\n    multioutput: {\"raw_values\", \"uniform_average\"} or array-like of shape (n_outputs,), default=\"uniform_average\"\n        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.\n\n        \"raw_values\" : Returns a full set of errors in case of multioutput input.\n\n        \"uniform_average\" : Errors of all outputs are averaged with uniform weight.\n\n    squared : bool, default=True\n        If True returns MSE value, if False returns RMSE (Root Mean Squared Error) value.\n\nReturns:\n    mse : mean squared error.\nExamples:\n\n    >>> mse_metric = datasets.load_metric(\"mse\")\n    >>> predictions = [2.5, 0.0, 2, 8]\n    >>> references = [3, -0.5, 2, 7]\n    >>> results = mse_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'mse': 0.375}\n    >>> rmse_result = mse_metric.compute(predictions=predictions, references=references, squared=False)\n    >>> print(rmse_result)\n    {'mse': 0.6123724356957945}\n\n    If you're using multi-dimensional lists, then set the config as follows :\n\n    >>> mse_metric = datasets.load_metric(\"mse\", \"multilist\")\n    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]\n    >>> references = [[0, 2], [-1, 2], [8, -5]]\n    >>> results = mse_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'mse': 0.7083333333333334}\n    >>> results = mse_metric.compute(predictions=predictions, references=references, multioutput='raw_values')\n    >>> print(results) # doctest: +NORMALIZE_WHITESPACE\n    {'mse': array([0.41666667, 1. ])}\n"


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



       def 				_a   (					self   )	->				Optional[int]:
           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 				_a   (					self   )	->				Tuple:
           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 				_a   (					self  ,UpperCAmelCase_  ,UpperCAmelCase_  ,UpperCAmelCase_=None  ,UpperCAmelCase_="uniform_average"  ,UpperCAmelCase_=True   )	->				Tuple:
           lowercase__									=    mean_squared_error(
               UpperCAmelCase_  ,UpperCAmelCase_  ,sample_weight=UpperCAmelCase_  ,multioutput=UpperCAmelCase_  ,squared=UpperCAmelCase_   )

           return {"mse": mse}




