


import gc
import unittest

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

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


UpperCAmelCase      =   False







class lowerCAmelCase_						(			unittest.TestCase	):

	'''simple docstring'''

	def 						UpperCamelCase__					(     self ):
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()



	@property
	def 						UpperCamelCase__					(     self ):
							return 12



	@property
	def 						UpperCamelCase__					(     self ):
							return 12



	@property
	def 						UpperCamelCase__					(     self ):
							return 32



	@property
	def 						UpperCamelCase__					(     self ):
							torch.manual_seed(0 )
							snake_case_	=				VQModel(
							    block_out_channels=[32, 64]		,				in_channels=3		,				out_channels=3		,				down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']		,				up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']		,				latent_channels=3		,				num_vq_embeddings=self.num_embed		,				vq_embed_dim=3		,				)
							return model



	@property
	def 						UpperCamelCase__					(     self ):
							snake_case_	=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )
							return tokenizer



	@property
	def 						UpperCamelCase__					(     self ):
							torch.manual_seed(0 )
							snake_case_	=				CLIPTextConfig(
							    bos_token_id=0		,				eos_token_id=2		,				hidden_size=self.text_embedder_hidden_size		,				intermediate_size=37		,				layer_norm_eps=1E-05		,				num_attention_heads=4		,				num_hidden_layers=5		,				pad_token_id=1		,				vocab_size=10_00		,				)
							return CLIPTextModel(_UpperCAmelCase )



	@property
	def 						UpperCamelCase__					(     self ):
							torch.manual_seed(0 )

							snake_case_	=				12
							snake_case_	=				12

							snake_case_	=				{
							    '''attention_bias''': True,
							    '''cross_attention_dim''': 32,
							    '''attention_head_dim''': height * width,
							    '''num_attention_heads''': 1,
							    '''num_vector_embeds''': self.num_embed,
							    '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
							    '''norm_num_groups''': 32,
							    '''sample_size''': width,
							    '''activation_fn''': '''geglu-approximate''',
							}

							snake_case_	=				TransformeraDModel(**_UpperCAmelCase )
							return model



	def 						UpperCamelCase__					(     self ):
							snake_case_	=				'''cpu'''

							snake_case_	=				self.dummy_vqvae
							snake_case_	=				self.dummy_text_encoder
							snake_case_	=				self.dummy_tokenizer
							snake_case_	=				self.dummy_transformer
							snake_case_	=				VQDiffusionScheduler(self.num_embed )
							snake_case_	=				LearnedClassifierFreeSamplingEmbeddings(learnable=_UpperCAmelCase )

							snake_case_	=				VQDiffusionPipeline(
							    vqvae=_UpperCAmelCase		,				text_encoder=_UpperCAmelCase		,				tokenizer=_UpperCAmelCase		,				transformer=_UpperCAmelCase		,				scheduler=_UpperCAmelCase		,				learned_classifier_free_sampling_embeddings=_UpperCAmelCase		,				)
							snake_case_	=				pipe.to(_UpperCAmelCase )
							pipe.set_progress_bar_config(disable=_UpperCAmelCase )

							snake_case_	=				'''teddy bear playing in the pool'''

							snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
							snake_case_	=				pipe([prompt]		,				generator=_UpperCAmelCase		,				num_inference_steps=2		,				output_type='''np''' )
							snake_case_	=				output.images

							snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
							snake_case_	=				pipe(
							    [prompt]		,				generator=_UpperCAmelCase		,				output_type='''np'''		,				return_dict=_UpperCAmelCase		,				num_inference_steps=2 )[0]

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

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

							snake_case_	=				np.array([0.6_551, 0.6_168, 0.5_008, 0.5_676, 0.5_659, 0.4_295, 0.6_073, 0.5_599, 0.4_992] )

							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 						UpperCamelCase__					(     self ):
							snake_case_	=				'''cpu'''

							snake_case_	=				self.dummy_vqvae
							snake_case_	=				self.dummy_text_encoder
							snake_case_	=				self.dummy_tokenizer
							snake_case_	=				self.dummy_transformer
							snake_case_	=				VQDiffusionScheduler(self.num_embed )
							snake_case_	=				LearnedClassifierFreeSamplingEmbeddings(
							    learnable=_UpperCAmelCase		,				hidden_size=self.text_embedder_hidden_size		,				length=tokenizer.model_max_length )

							snake_case_	=				VQDiffusionPipeline(
							    vqvae=_UpperCAmelCase		,				text_encoder=_UpperCAmelCase		,				tokenizer=_UpperCAmelCase		,				transformer=_UpperCAmelCase		,				scheduler=_UpperCAmelCase		,				learned_classifier_free_sampling_embeddings=_UpperCAmelCase		,				)
							snake_case_	=				pipe.to(_UpperCAmelCase )
							pipe.set_progress_bar_config(disable=_UpperCAmelCase )

							snake_case_	=				'''teddy bear playing in the pool'''

							snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
							snake_case_	=				pipe([prompt]		,				generator=_UpperCAmelCase		,				num_inference_steps=2		,				output_type='''np''' )
							snake_case_	=				output.images

							snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
							snake_case_	=				pipe(
							    [prompt]		,				generator=_UpperCAmelCase		,				output_type='''np'''		,				return_dict=_UpperCAmelCase		,				num_inference_steps=2 )[0]

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

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

							snake_case_	=				np.array([0.6_693, 0.6_075, 0.4_959, 0.5_701, 0.5_583, 0.4_333, 0.6_171, 0.5_684, 0.4_988] )

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







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

	'''simple docstring'''

	def 						UpperCamelCase__					(     self ):
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()





	def 						UpperCamelCase__					(     self ):
							snake_case_	=				load_numpy(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
							    '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy''' )

							snake_case_	=				VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq''' )
							snake_case_	=				pipeline.to(_UpperCAmelCase )
							pipeline.set_progress_bar_config(disable=_UpperCAmelCase )

							# requires GPU generator for gumbel softmax
							# don't use GPU generator in tests though
							snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(0 )
							snake_case_	=				pipeline(
							    '''teddy bear playing in the pool'''		,				num_images_per_prompt=1		,				generator=_UpperCAmelCase		,				output_type='''np'''		,				)

							snake_case_	=				output.images[0]

							assert image.shape == (2_56, 2_56, 3)
							assert np.abs(expected_image - image ).max() < 2.0



import numpy as np

import datasets


UpperCAmelCase      =   """
Compute the Mahalanobis Distance

Mahalonobis distance is the distance between a point and a distribution.
And not between two distinct points. It is effectively a multivariate equivalent of the Euclidean distance.
It was introduced by Prof. P. C. Mahalanobis in 1936
and has been used in various statistical applications ever since
[source: https://www.machinelearningplus.com/statistics/mahalanobis-distance/]
"""

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

UpperCAmelCase      =   """
Args:
    X: List of datapoints to be compared with the `reference_distribution`.
    reference_distribution: List of datapoints from the reference distribution we want to compare to.
Returns:
    mahalanobis: The Mahalonobis distance for each datapoint in `X`.
Examples:

    >>> mahalanobis_metric = datasets.load_metric(\"mahalanobis\")
    >>> results = mahalanobis_metric.compute(reference_distribution=[[0, 1], [1, 0]], X=[[0, 1]])
    >>> print(results)
    {'mahalanobis': array([0.5])}
"""







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

 '''simple docstring'''

 def 						UpperCamelCase__					(     self ):
       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 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       # convert to numpy arrays
       snake_case_	=				np.array(_UpperCAmelCase )
       snake_case_	=				np.array(_UpperCAmelCase )

       # 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
       snake_case_	=				X - np.mean(_UpperCAmelCase )
       snake_case_	=				np.cov(reference_distribution.T )
       try:
             snake_case_	=				np.linalg.inv(_UpperCAmelCase )
       except np.linalg.LinAlgError:
             snake_case_	=				np.linalg.pinv(_UpperCAmelCase )
       snake_case_	=				np.dot(_UpperCAmelCase		,				_UpperCAmelCase )
       snake_case_	=				np.dot(_UpperCAmelCase		,				X_minus_mu.T ).diagonal()

       return {"mahalanobis": mahal_dist}