
"""simple docstring"""


import gc
import unittest

import numpy as np
import torch
from torch.backends.cuda import sdp_kernel

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)
from diffusers.utils import randn_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_a, require_torch_gpu

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class 					lowerCAmelCase_					(   _lowercase		,       unittest.TestCase     ):



						'''simple docstring'''


						_lowerCamelCase:     List[str]   =     ConsistencyModelPipeline
						_lowerCamelCase:     Optional[int]   =     UNCONDITIONAL_IMAGE_GENERATION_PARAMS
						_lowerCamelCase:     Any   =     UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

						# Override required_optional_params to remove num_images_per_prompt
						_lowerCamelCase:     Any   =     frozenset(
						    [
						        '''num_inference_steps''',
						        '''generator''',
						        '''latents''',
						        '''output_type''',
						        '''return_dict''',
						        '''callback''',
						        '''callback_steps''',
						    ]     )


						@property
						def     _SCREAMING_SNAKE_CASE   (			self  :   Any   )  -> List[Any]:
							A     = UNetaDModel.from_pretrained(
							    'diffusers/consistency-models-test'	,subfolder='test_unet'	,)
							return unet


						@property
						def     _SCREAMING_SNAKE_CASE   (			self  :   str   )  -> Optional[Any]:
							A     = UNetaDModel.from_pretrained(
							    'diffusers/consistency-models-test'	,subfolder='test_unet_class_cond'	,)
							return unet


						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[Any]	,A_  :   List[Any]=False   )  -> Tuple:
							if class_cond:
								A     = self.dummy_cond_unet
							else:
								A     = self.dummy_uncond_unet

							# Default to CM multistep sampler
							A     = CMStochasticIterativeScheduler(
							    num_train_timesteps=40	,sigma_min=0.0_02	,sigma_max=80.0	,)

							A     = {
							    'unet': unet,
							    'scheduler': scheduler,
							}

							return components


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[Any]	,A_  :   Union[str, Any]	,A_  :   Dict=0   )  -> int:
							if str(A_   ).startswith('mps'   ):
								A     = torch.manual_seed(A_   )
							else:
								A     = torch.Generator(device=A_   ).manual_seed(A_   )

							A     = {
							    'batch_size': 1,
							    'num_inference_steps': None,
							    'timesteps': [22, 0],
							    'generator': generator,
							    'output_type': 'np',
							}

							return inputs


						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[int]   )  -> Dict:
							A     = 'cpu'  # ensure determinism for the device-dependent torch.Generator
							A     = self.get_dummy_components()
							A     = ConsistencyModelPipeline(**A_   )
							A     = pipe.to(A_   )
							pipe.set_progress_bar_config(disable=A_   )

							A     = self.get_dummy_inputs(A_   )
							A     = pipe(**A_   ).images
							assert image.shape == (1, 32, 32, 3)

							A     = image[0, -3:, -3:, -1]
							A     = np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]   )

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


						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[int]   )  -> Dict:
							A     = 'cpu'  # ensure determinism for the device-dependent torch.Generator
							A     = self.get_dummy_components(class_cond=A_   )
							A     = ConsistencyModelPipeline(**A_   )
							A     = pipe.to(A_   )
							pipe.set_progress_bar_config(disable=A_   )

							A     = self.get_dummy_inputs(A_   )
							A     = 0
							A     = pipe(**A_   ).images
							assert image.shape == (1, 32, 32, 3)

							A     = image[0, -3:, -3:, -1]
							A     = np.array([0.35_72, 0.62_73, 0.40_31, 0.39_61, 0.43_21, 0.57_30, 0.52_66, 0.47_80, 0.50_04]   )

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


						def     _SCREAMING_SNAKE_CASE   (			self  :   Dict   )  -> Optional[Any]:
							A     = 'cpu'  # ensure determinism for the device-dependent torch.Generator
							A     = self.get_dummy_components()
							A     = ConsistencyModelPipeline(**A_   )
							A     = pipe.to(A_   )
							pipe.set_progress_bar_config(disable=A_   )

							A     = self.get_dummy_inputs(A_   )
							A     = 1
							A     = None
							A     = pipe(**A_   ).images
							assert image.shape == (1, 32, 32, 3)

							A     = image[0, -3:, -3:, -1]
							A     = np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]   )

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




						def     _SCREAMING_SNAKE_CASE   (			self  :   Dict   )  -> Any:
							A     = 'cpu'  # ensure determinism for the device-dependent torch.Generator
							A     = self.get_dummy_components(class_cond=A_   )
							A     = ConsistencyModelPipeline(**A_   )
							A     = pipe.to(A_   )
							pipe.set_progress_bar_config(disable=A_   )

							A     = self.get_dummy_inputs(A_   )
							A     = 1
							A     = None
							A     = 0
							A     = pipe(**A_   ).images
							assert image.shape == (1, 32, 32, 3)

							A     = image[0, -3:, -3:, -1]
							A     = np.array([0.50_04, 0.50_04, 0.49_94, 0.50_08, 0.49_76, 0.50_18, 0.49_90, 0.49_82, 0.49_87]   )

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







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



						'''simple docstring'''


						def     _SCREAMING_SNAKE_CASE   (			self  :   Any   )  -> str:
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()


						def     _SCREAMING_SNAKE_CASE   (			self  :   Union[str, Any]	,A_  :   Any=0	,A_  :   Any=False	,A_  :   List[Any]="cpu"	,A_  :   str=torch.floataa	,A_  :   Any=(1, 3, 64, 64)   )  -> Tuple:
							A     = torch.manual_seed(A_   )

							A     = {
							    'num_inference_steps': None,
							    'timesteps': [22, 0],
							    'class_labels': 0,
							    'generator': generator,
							    'output_type': 'np',
							}

							if get_fixed_latents:
								A     = self.get_fixed_latents(seed=A_	,device=A_	,dtype=A_	,shape=A_   )
								A     = latents

							return inputs


						def     _SCREAMING_SNAKE_CASE   (			self  :   int	,A_  :   Dict=0	,A_  :   Any="cpu"	,A_  :   List[Any]=torch.floataa	,A_  :   List[str]=(1, 3, 64, 64)   )  -> List[Any]:
							if type(A_   ) == str:
								A     = torch.device(A_   )
							A     = torch.Generator(device=A_   ).manual_seed(A_   )
							A     = randn_tensor(A_	,generator=A_	,device=A_	,dtype=A_   )
							return latents


						def     _SCREAMING_SNAKE_CASE   (			self  :   List[str]   )  -> List[str]:
							A     = UNetaDModel.from_pretrained('diffusers/consistency_models'	,subfolder='diffusers_cd_imagenet64_l2'   )
							A     = CMStochasticIterativeScheduler(
							    num_train_timesteps=40	,sigma_min=0.0_02	,sigma_max=80.0	,)
							A     = ConsistencyModelPipeline(unet=A_	,scheduler=A_   )
							pipe.to(torch_device=A_   )
							pipe.set_progress_bar_config(disable=A_   )

							A     = self.get_inputs()
							A     = pipe(**A_   ).images
							assert image.shape == (1, 64, 64, 3)

							A     = image[0, -3:, -3:, -1]

							A     = np.array([0.08_88, 0.08_81, 0.06_66, 0.04_79, 0.02_92, 0.01_95, 0.02_01, 0.01_63, 0.02_54]   )

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


						def     _SCREAMING_SNAKE_CASE   (			self  :   Union[str, Any]   )  -> List[str]:
							A     = UNetaDModel.from_pretrained('diffusers/consistency_models'	,subfolder='diffusers_cd_imagenet64_l2'   )
							A     = CMStochasticIterativeScheduler(
							    num_train_timesteps=40	,sigma_min=0.0_02	,sigma_max=80.0	,)
							A     = ConsistencyModelPipeline(unet=A_	,scheduler=A_   )
							pipe.to(torch_device=A_   )
							pipe.set_progress_bar_config(disable=A_   )

							A     = self.get_inputs()
							A     = 1
							A     = None
							A     = pipe(**A_   ).images
							assert image.shape == (1, 64, 64, 3)

							A     = image[0, -3:, -3:, -1]

							A     = np.array([0.03_40, 0.01_52, 0.00_63, 0.02_67, 0.02_21, 0.01_07, 0.04_16, 0.01_86, 0.02_17]   )

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


						@require_torch_a
						def     _SCREAMING_SNAKE_CASE   (			self  :   Optional[int]   )  -> List[str]:
							A     = UNetaDModel.from_pretrained('diffusers/consistency_models'	,subfolder='diffusers_cd_imagenet64_l2'   )
							A     = CMStochasticIterativeScheduler(
							    num_train_timesteps=40	,sigma_min=0.0_02	,sigma_max=80.0	,)
							A     = ConsistencyModelPipeline(unet=A_	,scheduler=A_   )
							pipe.to(torch_device=A_	,torch_dtype=torch.floataa   )
							pipe.set_progress_bar_config(disable=A_   )

							A     = self.get_inputs(get_fixed_latents=A_	,device=A_   )
							# Ensure usage of flash attention in torch 2.0
							with sdp_kernel(enable_flash=A_	,enable_math=A_	,enable_mem_efficient=A_   ):
								A     = pipe(**A_   ).images
							assert image.shape == (1, 64, 64, 3)

							A     = image[0, -3:, -3:, -1]

							A     = np.array([0.18_75, 0.14_28, 0.12_89, 0.21_51, 0.20_92, 0.14_77, 0.18_77, 0.16_41, 0.13_53]   )

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




						@require_torch_a
						def     _SCREAMING_SNAKE_CASE   (			self  :   Dict   )  -> List[Any]:
							A     = UNetaDModel.from_pretrained('diffusers/consistency_models'	,subfolder='diffusers_cd_imagenet64_l2'   )
							A     = CMStochasticIterativeScheduler(
							    num_train_timesteps=40	,sigma_min=0.0_02	,sigma_max=80.0	,)
							A     = ConsistencyModelPipeline(unet=A_	,scheduler=A_   )
							pipe.to(torch_device=A_	,torch_dtype=torch.floataa   )
							pipe.set_progress_bar_config(disable=A_   )

							A     = self.get_inputs(get_fixed_latents=A_	,device=A_   )
							A     = 1
							A     = None
							# Ensure usage of flash attention in torch 2.0
							with sdp_kernel(enable_flash=A_	,enable_math=A_	,enable_mem_efficient=A_   ):
								A     = pipe(**A_   ).images
							assert image.shape == (1, 64, 64, 3)

							A     = image[0, -3:, -3:, -1]

							A     = np.array([0.16_63, 0.19_48, 0.22_75, 0.16_80, 0.12_04, 0.12_45, 0.18_58, 0.13_38, 0.20_95]   )

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







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}




