






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, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
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 ):

     '''simple docstring'''






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






           """simple docstring"""




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







     @property
     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           UpperCamelCase					       =  1
           UpperCamelCase					       =  3
           UpperCamelCase					       =  (32, 32)

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







     @property
     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           torch.manual_seed(0 )
           UpperCamelCase					       =  UNetaDConditionModel(
               block_out_channels=(32, 32, 64)   ,						layers_per_block=2   ,						sample_size=32   ,						in_channels=7   ,						out_channels=4   ,						down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D', 'CrossAttnDownBlock2D')   ,						up_block_types=('CrossAttnUpBlock2D', 'CrossAttnUpBlock2D', 'UpBlock2D')   ,						cross_attention_dim=32   ,						attention_head_dim=8   ,						use_linear_projection=_SCREAMING_SNAKE_CASE   ,						only_cross_attention=(True, True, False)   ,						num_class_embeds=100   ,						)
           return model







     @property
     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple ):






           """simple docstring"""




           torch.manual_seed(0 )
           UpperCamelCase					       =  AutoencoderKL(
               block_out_channels=[32, 32, 64]   ,						in_channels=3   ,						out_channels=3   ,						down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D', 'DownEncoderBlock2D']   ,						up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D']   ,						latent_channels=4   ,						)
           return model







     @property
     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any] ):






           """simple docstring"""




           torch.manual_seed(0 )
           UpperCamelCase					       =  CLIPTextConfig(
               bos_token_id=0   ,						eos_token_id=2   ,						hidden_size=32   ,						intermediate_size=37   ,						layer_norm_eps=1E-0_5   ,						num_attention_heads=4   ,						num_hidden_layers=5   ,						pad_token_id=1   ,						vocab_size=1000   ,						hidden_act='gelu'   ,						projection_dim=512   ,						)
           return CLIPTextModel(_SCREAMING_SNAKE_CASE )







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str] ):






           """simple docstring"""




           UpperCamelCase					       =  'cpu'  # ensure determinism for the device-dependent torch.Generator
           UpperCamelCase					       =  self.dummy_cond_unet_upscale
           UpperCamelCase					       =  DDPMScheduler()
           UpperCamelCase					       =  DDIMScheduler(prediction_type='v_prediction' )
           UpperCamelCase					       =  self.dummy_vae
           UpperCamelCase					       =  self.dummy_text_encoder
           UpperCamelCase					       =  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

           UpperCamelCase					       =  self.dummy_image.cpu().permute(0   ,						2   ,						3   ,						1 )[0]
           UpperCamelCase					       =  Image.fromarray(np.uinta(_SCREAMING_SNAKE_CASE ) ).convert('RGB' ).resize((64, 64) )

           # make sure here that pndm scheduler skips prk
           UpperCamelCase					       =  StableDiffusionUpscalePipeline(
               unet=_SCREAMING_SNAKE_CASE   ,						low_res_scheduler=_SCREAMING_SNAKE_CASE   ,						scheduler=_SCREAMING_SNAKE_CASE   ,						vae=_SCREAMING_SNAKE_CASE   ,						text_encoder=_SCREAMING_SNAKE_CASE   ,						tokenizer=_SCREAMING_SNAKE_CASE   ,						max_noise_level=350   ,						)
           UpperCamelCase					       =  sd_pipe.to(_SCREAMING_SNAKE_CASE )
           sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  'A painting of a squirrel eating a burger'
           UpperCamelCase					       =  torch.Generator(device=_SCREAMING_SNAKE_CASE ).manual_seed(0 )
           UpperCamelCase					       =  sd_pipe(
               [prompt]   ,						image=_SCREAMING_SNAKE_CASE   ,						generator=_SCREAMING_SNAKE_CASE   ,						guidance_scale=6.0   ,						noise_level=20   ,						num_inference_steps=2   ,						output_type='np'   ,						)

           UpperCamelCase					       =  output.images

           UpperCamelCase					       =  torch.Generator(device=_SCREAMING_SNAKE_CASE ).manual_seed(0 )
           UpperCamelCase					       =  sd_pipe(
               [prompt]   ,						image=_SCREAMING_SNAKE_CASE   ,						generator=_SCREAMING_SNAKE_CASE   ,						guidance_scale=6.0   ,						noise_level=20   ,						num_inference_steps=2   ,						output_type='np'   ,						return_dict=_SCREAMING_SNAKE_CASE   ,						)[0]

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

           UpperCamelCase					       =  low_res_image.size[0] * 4
           assert image.shape == (1, expected_height_width, expected_height_width, 3)
           UpperCamelCase					       =  np.array([0.3_1_1_3, 0.3_9_1_0, 0.4_2_7_2, 0.4_8_5_9, 0.5_0_6_1, 0.4_6_5_2, 0.5_3_6_2, 0.5_7_1_5, 0.5_6_6_1] )

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







     def   _SCREAMING_SNAKE_CASE     ( self      :    Any ):






           """simple docstring"""




           UpperCamelCase					       =  'cpu'  # ensure determinism for the device-dependent torch.Generator
           UpperCamelCase					       =  self.dummy_cond_unet_upscale
           UpperCamelCase					       =  DDPMScheduler()
           UpperCamelCase					       =  DDIMScheduler(prediction_type='v_prediction' )
           UpperCamelCase					       =  self.dummy_vae
           UpperCamelCase					       =  self.dummy_text_encoder
           UpperCamelCase					       =  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

           UpperCamelCase					       =  self.dummy_image.cpu().permute(0   ,						2   ,						3   ,						1 )[0]
           UpperCamelCase					       =  Image.fromarray(np.uinta(_SCREAMING_SNAKE_CASE ) ).convert('RGB' ).resize((64, 64) )

           # make sure here that pndm scheduler skips prk
           UpperCamelCase					       =  StableDiffusionUpscalePipeline(
               unet=_SCREAMING_SNAKE_CASE   ,						low_res_scheduler=_SCREAMING_SNAKE_CASE   ,						scheduler=_SCREAMING_SNAKE_CASE   ,						vae=_SCREAMING_SNAKE_CASE   ,						text_encoder=_SCREAMING_SNAKE_CASE   ,						tokenizer=_SCREAMING_SNAKE_CASE   ,						max_noise_level=350   ,						)
           UpperCamelCase					       =  sd_pipe.to(_SCREAMING_SNAKE_CASE )
           sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  'A painting of a squirrel eating a burger'
           UpperCamelCase					       =  sd_pipe(
               2 * [prompt]   ,						image=2 * [low_res_image]   ,						guidance_scale=6.0   ,						noise_level=20   ,						num_inference_steps=2   ,						output_type='np'   ,						)
           UpperCamelCase					       =  output.images
           assert image.shape[0] == 2

           UpperCamelCase					       =  torch.Generator(device=_SCREAMING_SNAKE_CASE ).manual_seed(0 )
           UpperCamelCase					       =  sd_pipe(
               [prompt]   ,						image=_SCREAMING_SNAKE_CASE   ,						generator=_SCREAMING_SNAKE_CASE   ,						num_images_per_prompt=2   ,						guidance_scale=6.0   ,						noise_level=20   ,						num_inference_steps=2   ,						output_type='np'   ,						)
           UpperCamelCase					       =  output.images
           assert image.shape[0] == 2







     @unittest.skipIf(torch_device != 'cuda'   ,						'This test requires a GPU' )
     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  self.dummy_cond_unet_upscale
           UpperCamelCase					       =  DDPMScheduler()
           UpperCamelCase					       =  DDIMScheduler(prediction_type='v_prediction' )
           UpperCamelCase					       =  self.dummy_vae
           UpperCamelCase					       =  self.dummy_text_encoder
           UpperCamelCase					       =  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

           UpperCamelCase					       =  self.dummy_image.cpu().permute(0   ,						2   ,						3   ,						1 )[0]
           UpperCamelCase					       =  Image.fromarray(np.uinta(_SCREAMING_SNAKE_CASE ) ).convert('RGB' ).resize((64, 64) )

           # put models in fp16, except vae as it overflows in fp16
           UpperCamelCase					       =  unet.half()
           UpperCamelCase					       =  text_encoder.half()

           # make sure here that pndm scheduler skips prk
           UpperCamelCase					       =  StableDiffusionUpscalePipeline(
               unet=_SCREAMING_SNAKE_CASE   ,						low_res_scheduler=_SCREAMING_SNAKE_CASE   ,						scheduler=_SCREAMING_SNAKE_CASE   ,						vae=_SCREAMING_SNAKE_CASE   ,						text_encoder=_SCREAMING_SNAKE_CASE   ,						tokenizer=_SCREAMING_SNAKE_CASE   ,						max_noise_level=350   ,						)
           UpperCamelCase					       =  sd_pipe.to(_SCREAMING_SNAKE_CASE )
           sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )

           UpperCamelCase					       =  'A painting of a squirrel eating a burger'
           UpperCamelCase					       =  torch.manual_seed(0 )
           UpperCamelCase					       =  sd_pipe(
               [prompt]   ,						image=_SCREAMING_SNAKE_CASE   ,						generator=_SCREAMING_SNAKE_CASE   ,						num_inference_steps=2   ,						output_type='np'   ,						).images

           UpperCamelCase					       =  low_res_image.size[0] * 4
           assert image.shape == (1, expected_height_width, expected_height_width, 3)





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

     '''simple docstring'''






     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str] ):






           """simple docstring"""




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







     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           UpperCamelCase					       =  load_image(
               'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
               '/sd2-upscale/low_res_cat.png' )
           UpperCamelCase					       =  load_numpy(
               'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'
               '/upsampled_cat.npy' )

           UpperCamelCase					       =  'stabilityai/stable-diffusion-x4-upscaler'
           UpperCamelCase					       =  StableDiffusionUpscalePipeline.from_pretrained(_SCREAMING_SNAKE_CASE )
           pipe.to(_SCREAMING_SNAKE_CASE )
           pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )
           pipe.enable_attention_slicing()

           UpperCamelCase					       =  'a cat sitting on a park bench'

           UpperCamelCase					       =  torch.manual_seed(0 )
           UpperCamelCase					       =  pipe(
               prompt=_SCREAMING_SNAKE_CASE   ,						image=_SCREAMING_SNAKE_CASE   ,						generator=_SCREAMING_SNAKE_CASE   ,						output_type='np'   ,						)
           UpperCamelCase					       =  output.images[0]

           assert image.shape == (512, 512, 3)
           assert np.abs(expected_image - image ).max() < 1E-3







     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           UpperCamelCase					       =  load_image(
               'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
               '/sd2-upscale/low_res_cat.png' )
           UpperCamelCase					       =  load_numpy(
               'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'
               '/upsampled_cat_fp16.npy' )

           UpperCamelCase					       =  'stabilityai/stable-diffusion-x4-upscaler'
           UpperCamelCase					       =  StableDiffusionUpscalePipeline.from_pretrained(
               _SCREAMING_SNAKE_CASE   ,						torch_dtype=torch.floataa   ,						)
           pipe.to(_SCREAMING_SNAKE_CASE )
           pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )
           pipe.enable_attention_slicing()

           UpperCamelCase					       =  'a cat sitting on a park bench'

           UpperCamelCase					       =  torch.manual_seed(0 )
           UpperCamelCase					       =  pipe(
               prompt=_SCREAMING_SNAKE_CASE   ,						image=_SCREAMING_SNAKE_CASE   ,						generator=_SCREAMING_SNAKE_CASE   ,						output_type='np'   ,						)
           UpperCamelCase					       =  output.images[0]

           assert image.shape == (512, 512, 3)
           assert np.abs(expected_image - image ).max() < 5E-1







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str] ):






           """simple docstring"""




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

           UpperCamelCase					       =  load_image(
               'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
               '/sd2-upscale/low_res_cat.png' )

           UpperCamelCase					       =  'stabilityai/stable-diffusion-x4-upscaler'
           UpperCamelCase					       =  StableDiffusionUpscalePipeline.from_pretrained(
               _SCREAMING_SNAKE_CASE   ,						torch_dtype=torch.floataa   ,						)
           pipe.to(_SCREAMING_SNAKE_CASE )
           pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE )
           pipe.enable_attention_slicing(1 )
           pipe.enable_sequential_cpu_offload()

           UpperCamelCase					       =  'a cat sitting on a park bench'

           UpperCamelCase					       =  torch.manual_seed(0 )
           UpperCamelCase					       =  pipe(
               prompt=_SCREAMING_SNAKE_CASE   ,						image=_SCREAMING_SNAKE_CASE   ,						generator=_SCREAMING_SNAKE_CASE   ,						num_inference_steps=5   ,						output_type='np'   ,						)

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











import numpy as np

import datasets


A__      :     int       =					'\nCompute the Mahalanobis Distance\n\nMahalonobis distance is the distance between a point and a distribution.\nAnd not between two distinct points. It is effectively a multivariate equivalent of the Euclidean distance.\nIt was introduced by Prof. P. C. Mahalanobis in 1936\nand has been used in various statistical applications ever since\n[source: https://www.machinelearningplus.com/statistics/mahalanobis-distance/]\n'

A__      :     Optional[int]       =					'\\n@article{de2000mahalanobis,\n  title={The mahalanobis distance},\n  author={De Maesschalck, Roy and Jouan-Rimbaud, Delphine and Massart, D{\'e}sir{\'e} L},\n  journal={Chemometrics and intelligent laboratory systems},\n  volume={50},\n  number={1},\n  pages={1--18},\n  year={2000},\n  publisher={Elsevier}\n}\n'

A__      :     Optional[int]       =					'\nArgs:\n    X: List of datapoints to be compared with the `reference_distribution`.\n    reference_distribution: List of datapoints from the reference distribution we want to compare to.\nReturns:\n    mahalanobis: The Mahalonobis distance for each datapoint in `X`.\nExamples:\n\n    >>> mahalanobis_metric = datasets.load_metric("mahalanobis")\n    >>> results = mahalanobis_metric.compute(reference_distribution=[[0, 1], [1, 0]], X=[[0, 1]])\n    >>> print(results)\n    {\'mahalanobis\': array([0.5])}\n'
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION						,_KWARGS_DESCRIPTION      )
class  _UpperCAmelCase   (			datasets.Metric      ):

				"""simple docstring"""




				def 							lowercase__  (	self		:       str    ):

								'''simple docstring'''
								return datasets.MetricInfo(
								    description=_DESCRIPTION,    citation=_CITATION,    inputs_description=_KWARGS_DESCRIPTION,    features=datasets.Features(
								        {
								            '''X''': datasets.Sequence(datasets.Value('''float''',    id='''sequence'''    ),    id='''X'''    ),
								        }    ),    )




				def 							lowercase__  (	self		:       Dict,    lowerCamelCase		:       int,    lowerCamelCase		:       Optional[int]    ):

								'''simple docstring'''
								# convert to numpy arrays
								lowercase__  					=     np.array(lowerCamelCase    )
								lowercase__  					=     np.array(lowerCamelCase    )

								# Assert that arrays are 2D
								if len(X.shape    ) != 2:
												raise ValueError('''Expected `X` to be a 2D vector'''    )
								if len(reference_distribution.shape    ) != 2:
												raise ValueError('''Expected `reference_distribution` to be a 2D vector'''    )
								if reference_distribution.shape[0] < 2:
												raise ValueError(
												    '''Expected `reference_distribution` to be a 2D vector with more than one element in the first dimension'''    )

								# Get mahalanobis distance for each prediction
								lowercase__  					=     X - np.mean(lowerCamelCase    )
								lowercase__  					=     np.cov(reference_distribution.T    )
								try:
												lowercase__  					=     np.linalg.inv(lowerCamelCase    )
								except np.linalg.LinAlgError:
												lowercase__  					=     np.linalg.pinv(lowerCamelCase    )
								lowercase__  					=     np.dot(lowerCamelCase,    lowerCamelCase    )
								lowercase__  					=     np.dot(lowerCamelCase,    X_minus_mu.T    ).diagonal()

								return {"mahalanobis": mahal_dist}




