




import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
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


enable_full_determinism()
class       a__   ( unittest.TestCase				):






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



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






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =       1
											__lowerCAmelCase           =       3
											__lowerCAmelCase           =       (3_2, 3_2)

											__lowerCAmelCase           =       floats_tensor((batch_size, num_channels) + sizes		,      rng=random.Random(0      )      ).to(UpperCamelCase__      )
											return image






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											torch.manual_seed(0      )
											__lowerCAmelCase           =       UNetaDConditionModel(
											    block_out_channels=(3_2, 6_4)		,      layers_per_block=2		,      sample_size=3_2		,      in_channels=4		,      out_channels=4		,      down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")		,      up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")		,      cross_attention_dim=3_2		,      )
											return model






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											torch.manual_seed(0      )
											__lowerCAmelCase           =       AutoencoderKL(
											    block_out_channels=[3_2, 6_4]		,      in_channels=3		,      out_channels=3		,      down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]		,      up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]		,      latent_channels=4		,      )
											return model






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											torch.manual_seed(0      )
											__lowerCAmelCase           =       RobertaSeriesConfig(
											    hidden_size=3_2		,      project_dim=3_2		,      intermediate_size=3_7		,      layer_norm_eps=1E-0_5		,      num_attention_heads=4		,      num_hidden_layers=5		,      pad_token_id=1		,      vocab_size=5_0_0_6		,      )
											return RobertaSeriesModelWithTransformation(UpperCamelCase__      )






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											def extract(*_A		,      **_A      ):
																class       a__   :






																						def __init__(       self      ):




																											"""simple docstring"""



																											__lowerCAmelCase           =       torch.ones([0]      )






																						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




																											"""simple docstring"""



																											self.pixel_values.to(UpperCamelCase__      )
																											return self






																return Out()

											return extract






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =       "cpu"  # ensure determinism for the device-dependent torch.Generator
											__lowerCAmelCase           =       self.dummy_cond_unet
											__lowerCAmelCase           =       PNDMScheduler(skip_prk_steps=UpperCamelCase__      )
											__lowerCAmelCase           =       self.dummy_vae
											__lowerCAmelCase           =       self.dummy_text_encoder
											__lowerCAmelCase           =       XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"      )
											__lowerCAmelCase           =       7_7

											__lowerCAmelCase           =       self.dummy_image.to(UpperCamelCase__      )
											__lowerCAmelCase           =       init_image / 2 + 0.5

											# make sure here that pndm scheduler skips prk
											__lowerCAmelCase           =       AltDiffusionImgaImgPipeline(
											    unet=UpperCamelCase__		,      scheduler=UpperCamelCase__		,      vae=UpperCamelCase__		,      text_encoder=UpperCamelCase__		,      tokenizer=UpperCamelCase__		,      safety_checker=UpperCamelCase__		,      feature_extractor=self.dummy_extractor		,      )
											__lowerCAmelCase           =       VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor		,      do_normalize=UpperCamelCase__      )
											__lowerCAmelCase           =       alt_pipe.to(UpperCamelCase__      )
											alt_pipe.set_progress_bar_config(disable=UpperCamelCase__      )

											__lowerCAmelCase           =       "A painting of a squirrel eating a burger"
											__lowerCAmelCase           =       torch.Generator(device=UpperCamelCase__      ).manual_seed(0      )
											__lowerCAmelCase           =       alt_pipe(
											    [prompt]		,      generator=UpperCamelCase__		,      guidance_scale=6.0		,      num_inference_steps=2		,      output_type="np"		,      image=UpperCamelCase__		,      )

											__lowerCAmelCase           =       output.images

											__lowerCAmelCase           =       torch.Generator(device=UpperCamelCase__      ).manual_seed(0      )
											__lowerCAmelCase           =       alt_pipe(
											    [prompt]		,      generator=UpperCamelCase__		,      guidance_scale=6.0		,      num_inference_steps=2		,      output_type="np"		,      image=UpperCamelCase__		,      return_dict=UpperCamelCase__		,      )[0]

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

											assert image.shape == (1, 3_2, 3_2, 3)
											__lowerCAmelCase           =       np.array([0.44_27, 0.37_31, 0.42_49, 0.49_41, 0.45_46, 0.41_48, 0.41_93, 0.46_66, 0.44_99]      )

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






						@unittest.skipIf(torch_device != "cuda"		,      "This test requires a GPU"      )
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =       self.dummy_cond_unet
											__lowerCAmelCase           =       PNDMScheduler(skip_prk_steps=UpperCamelCase__      )
											__lowerCAmelCase           =       self.dummy_vae
											__lowerCAmelCase           =       self.dummy_text_encoder
											__lowerCAmelCase           =       XLMRobertaTokenizer.from_pretrained("hf-internal-testing/tiny-xlm-roberta"      )
											__lowerCAmelCase           =       7_7

											__lowerCAmelCase           =       self.dummy_image.to(UpperCamelCase__      )

											# put models in fp16
											__lowerCAmelCase           =       unet.half()
											__lowerCAmelCase           =       vae.half()
											__lowerCAmelCase           =       bert.half()

											# make sure here that pndm scheduler skips prk
											__lowerCAmelCase           =       AltDiffusionImgaImgPipeline(
											    unet=UpperCamelCase__		,      scheduler=UpperCamelCase__		,      vae=UpperCamelCase__		,      text_encoder=UpperCamelCase__		,      tokenizer=UpperCamelCase__		,      safety_checker=UpperCamelCase__		,      feature_extractor=self.dummy_extractor		,      )
											__lowerCAmelCase           =       VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor		,      do_normalize=UpperCamelCase__      )
											__lowerCAmelCase           =       alt_pipe.to(UpperCamelCase__      )
											alt_pipe.set_progress_bar_config(disable=UpperCamelCase__      )

											__lowerCAmelCase           =       "A painting of a squirrel eating a burger"
											__lowerCAmelCase           =       torch.manual_seed(0      )
											__lowerCAmelCase           =       alt_pipe(
											    [prompt]		,      generator=UpperCamelCase__		,      num_inference_steps=2		,      output_type="np"		,      image=UpperCamelCase__		,      ).images

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






						@unittest.skipIf(torch_device != "cuda"		,      "This test requires a GPU"      )
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =       load_image(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    "/img2img/sketch-mountains-input.jpg"      )
											# resize to resolution that is divisible by 8 but not 16 or 32
											__lowerCAmelCase           =       init_image.resize((7_6_0, 5_0_4)      )

											__lowerCAmelCase           =       "BAAI/AltDiffusion"
											__lowerCAmelCase           =       AltDiffusionImgaImgPipeline.from_pretrained(
											    UpperCamelCase__		,      safety_checker=UpperCamelCase__		,      )
											pipe.to(UpperCamelCase__      )
											pipe.set_progress_bar_config(disable=UpperCamelCase__      )
											pipe.enable_attention_slicing()

											__lowerCAmelCase           =       "A fantasy landscape, trending on artstation"

											__lowerCAmelCase           =       torch.manual_seed(0      )
											__lowerCAmelCase           =       pipe(
											    prompt=UpperCamelCase__		,      image=UpperCamelCase__		,      strength=0.75		,      guidance_scale=7.5		,      generator=UpperCamelCase__		,      output_type="np"		,      )
											__lowerCAmelCase           =       output.images[0]

											__lowerCAmelCase           =       image[2_5_5:2_5_8, 3_8_3:3_8_6, -1]

											assert image.shape == (5_0_4, 7_6_0, 3)
											__lowerCAmelCase           =       np.array([0.93_58, 0.93_97, 0.95_99, 0.99_01, 1.00_00, 1.00_00, 0.98_82, 1.00_00, 1.00_00]      )

											assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-2






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






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



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






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =       load_image(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
											    "/img2img/sketch-mountains-input.jpg"      )
											__lowerCAmelCase           =       init_image.resize((7_6_8, 5_1_2)      )
											__lowerCAmelCase           =       load_numpy(
											    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy"      )

											__lowerCAmelCase           =       "BAAI/AltDiffusion"
											__lowerCAmelCase           =       AltDiffusionImgaImgPipeline.from_pretrained(
											    UpperCamelCase__		,      safety_checker=UpperCamelCase__		,      )
											pipe.to(UpperCamelCase__      )
											pipe.set_progress_bar_config(disable=UpperCamelCase__      )
											pipe.enable_attention_slicing()

											__lowerCAmelCase           =       "A fantasy landscape, trending on artstation"

											__lowerCAmelCase           =       torch.manual_seed(0      )
											__lowerCAmelCase           =       pipe(
											    prompt=UpperCamelCase__		,      image=UpperCamelCase__		,      strength=0.75		,      guidance_scale=7.5		,      generator=UpperCamelCase__		,      output_type="np"		,      )
											__lowerCAmelCase           =       output.images[0]

											assert image.shape == (5_1_2, 7_6_8, 3)
											# img2img is flaky across GPUs even in fp32, so using MAE here
											assert np.abs(expected_image - image      ).max() < 1E-2










from sklearn.metrics import mean_squared_error

import datasets


UpperCamelCase__         =      """\
@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}
}
"""

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


UpperCamelCase__         =      """
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       a__   ( datasets.Metric				):






						def     __SCREAMING_SNAKE_CASE(       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     __SCREAMING_SNAKE_CASE(       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     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A=None		,      _A="uniform_average"		,      _A=True      ):




											"""simple docstring"""



											__lowerCAmelCase           =	mean_squared_error(
											    _A		,      _A		,      sample_weight=_A		,      multioutput=_A		,      squared=_A      )

											return {"mse": mse}




