


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        _lowercase      (			A__						,			unittest.TestCase       ):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__				:			Dict						=		ConsistencyModelPipeline
				SCREAMING_SNAKE_CASE__				:			List[Any]						=		UNCONDITIONAL_IMAGE_GENERATION_PARAMS
				SCREAMING_SNAKE_CASE__				:			Optional[Any]						=		UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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


				@property
				def 		__magic_name__(       self			:Optional[int]	)		->			Optional[int]:
										__SCREAMING_SNAKE_CASE						:     Tuple					=		UNetaDModel.from_pretrained(
										    '''diffusers/consistency-models-test'''			,			subfolder='''test_unet'''			,			)
										return unet


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


				def 		__magic_name__(       self			:Dict			,			lowerCAmelCase__			:List[str]=False	)		->			Optional[int]:
										if class_cond:
																__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		self.dummy_cond_unet
										else:
																__SCREAMING_SNAKE_CASE						:     Dict					=		self.dummy_uncond_unet

										# Default to CM multistep sampler
										__SCREAMING_SNAKE_CASE						:     List[str]					=		CMStochasticIterativeScheduler(
										    num_train_timesteps=40			,			sigma_min=0.002			,			sigma_max=80.0			,			)

										__SCREAMING_SNAKE_CASE						:     Dict					=		{
										    '''unet''': unet,
										    '''scheduler''': scheduler,
										}

										return components


				def 		__magic_name__(       self			:int			,			lowerCAmelCase__			:List[str]			,			lowerCAmelCase__			:Optional[int]=0	)		->			Any:
										if str(lowerCAmelCase__	).startswith('''mps'''	):
																__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		torch.manual_seed(lowerCAmelCase__	)
										else:
																__SCREAMING_SNAKE_CASE						:     List[str]					=		torch.Generator(device=lowerCAmelCase__	).manual_seed(lowerCAmelCase__	)

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

										return inputs


				def 		__magic_name__(       self			:Any	)		->			str:
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__SCREAMING_SNAKE_CASE						:     List[str]					=		self.get_dummy_components()
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		ConsistencyModelPipeline(**lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		pipe.to(lowerCAmelCase__	)
										pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Dict					=		self.get_dummy_inputs(lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		pipe(**lowerCAmelCase__	).images
										assert image.shape == (1, 32, 32, 3)

										__SCREAMING_SNAKE_CASE						:     str					=		image[0, -3:, -3:, -1]
										__SCREAMING_SNAKE_CASE						:     Dict					=		np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]	)

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


				def 		__magic_name__(       self			:Optional[Any]	)		->			List[str]:
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__SCREAMING_SNAKE_CASE						:     Dict					=		self.get_dummy_components(class_cond=lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		ConsistencyModelPipeline(**lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     str					=		pipe.to(lowerCAmelCase__	)
										pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		self.get_dummy_inputs(lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		0
										__SCREAMING_SNAKE_CASE						:     Tuple					=		pipe(**lowerCAmelCase__	).images
										assert image.shape == (1, 32, 32, 3)

										__SCREAMING_SNAKE_CASE						:     Dict					=		image[0, -3:, -3:, -1]
										__SCREAMING_SNAKE_CASE						:     Dict					=		np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]	)

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


				def 		__magic_name__(       self			:str	)		->			Optional[int]:
										__SCREAMING_SNAKE_CASE						:     int					=		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__SCREAMING_SNAKE_CASE						:     List[str]					=		self.get_dummy_components()
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		ConsistencyModelPipeline(**lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     List[str]					=		pipe.to(lowerCAmelCase__	)
										pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     List[Any]					=		self.get_dummy_inputs(lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     Tuple					=		1
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		None
										__SCREAMING_SNAKE_CASE						:     Dict					=		pipe(**lowerCAmelCase__	).images
										assert image.shape == (1, 32, 32, 3)

										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		image[0, -3:, -3:, -1]
										__SCREAMING_SNAKE_CASE						:     int					=		np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]	)

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



				def 		__magic_name__(       self			:str	)		->			str:
										__SCREAMING_SNAKE_CASE						:     Dict					=		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
										__SCREAMING_SNAKE_CASE						:     Any					=		self.get_dummy_components(class_cond=lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		ConsistencyModelPipeline(**lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     Tuple					=		pipe.to(lowerCAmelCase__	)
										pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Any					=		self.get_dummy_inputs(lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     int					=		1
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		None
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		0
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		pipe(**lowerCAmelCase__	).images
										assert image.shape == (1, 32, 32, 3)

										__SCREAMING_SNAKE_CASE						:     Dict					=		image[0, -3:, -3:, -1]
										__SCREAMING_SNAKE_CASE						:     int					=		np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]	)

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




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






				'''simple docstring'''



				def 		__magic_name__(       self			:Union[str, Any]	)		->			Optional[Any]:
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()


				def 		__magic_name__(       self			:Dict			,			lowerCAmelCase__			:Tuple=0			,			lowerCAmelCase__			:Union[str, Any]=False			,			lowerCAmelCase__			:int="cpu"			,			lowerCAmelCase__			:Optional[Any]=torch.floataa			,			lowerCAmelCase__			:Union[str, Any]=(1, 3, 64, 64)	)		->			Optional[int]:
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		torch.manual_seed(lowerCAmelCase__	)

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

										if get_fixed_latents:
																__SCREAMING_SNAKE_CASE						:     str					=		self.get_fixed_latents(seed=lowerCAmelCase__			,			device=lowerCAmelCase__			,			dtype=lowerCAmelCase__			,			shape=lowerCAmelCase__	)
																__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		latents

										return inputs


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


				def 		__magic_name__(       self			:str	)		->			List[str]:
										__SCREAMING_SNAKE_CASE						:     Tuple					=		UNetaDModel.from_pretrained('''diffusers/consistency_models'''			,			subfolder='''diffusers_cd_imagenet64_l2'''	)
										__SCREAMING_SNAKE_CASE						:     str					=		CMStochasticIterativeScheduler(
										    num_train_timesteps=40			,			sigma_min=0.002			,			sigma_max=80.0			,			)
										__SCREAMING_SNAKE_CASE						:     List[str]					=		ConsistencyModelPipeline(unet=lowerCAmelCase__			,			scheduler=lowerCAmelCase__	)
										pipe.to(torch_device=lowerCAmelCase__	)
										pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     List[Any]					=		self.get_inputs()
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		pipe(**lowerCAmelCase__	).images
										assert image.shape == (1, 64, 64, 3)

										__SCREAMING_SNAKE_CASE						:     Tuple					=		image[0, -3:, -3:, -1]

										__SCREAMING_SNAKE_CASE						:     Dict					=		np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254]	)

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


				def 		__magic_name__(       self			:List[Any]	)		->			Optional[int]:
										__SCREAMING_SNAKE_CASE						:     List[str]					=		UNetaDModel.from_pretrained('''diffusers/consistency_models'''			,			subfolder='''diffusers_cd_imagenet64_l2'''	)
										__SCREAMING_SNAKE_CASE						:     Dict					=		CMStochasticIterativeScheduler(
										    num_train_timesteps=40			,			sigma_min=0.002			,			sigma_max=80.0			,			)
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		ConsistencyModelPipeline(unet=lowerCAmelCase__			,			scheduler=lowerCAmelCase__	)
										pipe.to(torch_device=lowerCAmelCase__	)
										pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Tuple					=		self.get_inputs()
										__SCREAMING_SNAKE_CASE						:     str					=		1
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		None
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		pipe(**lowerCAmelCase__	).images
										assert image.shape == (1, 64, 64, 3)

										__SCREAMING_SNAKE_CASE						:     int					=		image[0, -3:, -3:, -1]

										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217]	)

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


				@require_torch_a
				def 		__magic_name__(       self			:Union[str, Any]	)		->			int:
										__SCREAMING_SNAKE_CASE						:     Dict					=		UNetaDModel.from_pretrained('''diffusers/consistency_models'''			,			subfolder='''diffusers_cd_imagenet64_l2'''	)
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		CMStochasticIterativeScheduler(
										    num_train_timesteps=40			,			sigma_min=0.002			,			sigma_max=80.0			,			)
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		ConsistencyModelPipeline(unet=lowerCAmelCase__			,			scheduler=lowerCAmelCase__	)
										pipe.to(torch_device=lowerCAmelCase__			,			torch_dtype=torch.floataa	)
										pipe.set_progress_bar_config(disable=lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     str					=		self.get_inputs(get_fixed_latents=lowerCAmelCase__			,			device=lowerCAmelCase__	)
										# Ensure usage of flash attention in torch 2.0
										with sdp_kernel(enable_flash=lowerCAmelCase__			,			enable_math=lowerCAmelCase__			,			enable_mem_efficient=lowerCAmelCase__	):
																__SCREAMING_SNAKE_CASE						:     List[Any]					=		pipe(**lowerCAmelCase__	).images
										assert image.shape == (1, 64, 64, 3)

										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		image[0, -3:, -3:, -1]

										__SCREAMING_SNAKE_CASE						:     List[Any]					=		np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353]	)

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



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

										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		self.get_inputs(get_fixed_latents=lowerCAmelCase__			,			device=lowerCAmelCase__	)
										__SCREAMING_SNAKE_CASE						:     List[Any]					=		1
										__SCREAMING_SNAKE_CASE						:     Any					=		None
										# Ensure usage of flash attention in torch 2.0
										with sdp_kernel(enable_flash=lowerCAmelCase__			,			enable_math=lowerCAmelCase__			,			enable_mem_efficient=lowerCAmelCase__	):
																__SCREAMING_SNAKE_CASE						:     List[str]					=		pipe(**lowerCAmelCase__	).images
										assert image.shape == (1, 64, 64, 3)

										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		image[0, -3:, -3:, -1]

										__SCREAMING_SNAKE_CASE						:     Any					=		np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095]	)

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




import math_equivalence  # From: git+https://github.com/hendrycks/math.git

import datasets


__lowerCAmelCase	: Optional[Any]  				='\\n@article{hendrycksmath2021,\n  title={Measuring Mathematical Problem Solving With the MATH Dataset},\n  author={Dan Hendrycks\n    and Collin Burns\n    and Saurav Kadavath\n    and Akul Arora\n    and Steven Basart\n    and Eric Tang\n    and Dawn Song\n    and Jacob Steinhardt},\n  journal={arXiv preprint arXiv:2103.03874},\n  year={2021}\n}\n'


__lowerCAmelCase	: Any  				='\\nThis metric is used to assess performance on the Mathematics Aptitude Test of Heuristics (MATH) dataset.\nIt first canonicalizes the inputs (e.g., converting "1/2" to "\\frac{1}{2}") and then computes accuracy.\n'


__lowerCAmelCase	: Optional[Any]  				=r'\nCalculates accuracy after canonicalizing inputs.\n\nArgs:\n    predictions: list of predictions to score. Each prediction\n        is a string that contains natural language and LaTex.\n    references: list of reference for each prediction. Each\n        reference is a string that contains natural language\n        and LaTex.\nReturns:\n    accuracy: accuracy after canonicalizing inputs\n        (e.g., converting "1/2" to "\\frac{1}{2}")\n\nExamples:\n    >>> metric = datasets.load_metric("competition_math")\n    >>> results = metric.compute(references=["\\frac{1}{2}"], predictions=["1/2"])\n    >>> print(results)\n    {\'accuracy\': 1.0}\n'
@datasets.utils.file_utils.add_end_docstrings(_DESCRIPTION						,			_KWARGS_DESCRIPTION       )
class        _lowercase      (			datasets.Metric       ):






				'''simple docstring'''



				def 		__magic_name__(       self			:Optional[Any]	)		->			List[str]:
										return datasets.MetricInfo(
										    description=_DESCRIPTION			,			citation=_CITATION			,			inputs_description=_KWARGS_DESCRIPTION			,			features=datasets.Features(
										        {
										            '''predictions''': datasets.Value('''string'''	),
										            '''references''': datasets.Value('''string'''	),
										        }	)			,			homepage='''https://github.com/hendrycks/math'''			,			codebase_urls=['''https://github.com/hendrycks/math''']			,			)



				def 		__magic_name__(       self			:Any			,			lowerCAmelCase__			:Optional[int]			,			lowerCAmelCase__			:Any	)		->			Optional[int]:
										__SCREAMING_SNAKE_CASE						:     Tuple					=		0.0
										for i, j in zip(lowerCAmelCase__			,			lowerCAmelCase__	):
																n_correct += 1.0 if math_equivalence.is_equiv(lowerCAmelCase__			,			lowerCAmelCase__	) else 0.0
										__SCREAMING_SNAKE_CASE						:     str					=		n_correct / len(lowerCAmelCase__	)
										return {
										    "accuracy": accuracy,
										}
