"""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_			(	UpperCamelCase__   ,			unittest.TestCase    ):



				"""simple docstring"""
				_lowerCAmelCase   :			Any	     =  ConsistencyModelPipeline
				_lowerCAmelCase   :			int	     =  UNCONDITIONAL_IMAGE_GENERATION_PARAMS
				_lowerCAmelCase   :			str	     =  UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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

				@property
				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	UNetaDModel.from_pretrained(
										    'diffusers/consistency-models-test'						,			subfolder='test_unet'						,			)
										return unet

				@property
				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	UNetaDModel.from_pretrained(
										    'diffusers/consistency-models-test'						,			subfolder='test_unet_class_cond'						,			)
										return unet

				def 						snake_case       ( self						,			lowerCAmelCase=False			):






										"""simple docstring"""

										if class_cond:
																snake_case	    =	self.dummy_cond_unet
										else:
																snake_case	    =	self.dummy_uncond_unet

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

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

										return components

				def 						snake_case       ( self						,			lowerCAmelCase						,			lowerCAmelCase=0			):






										"""simple docstring"""

										if str(__a			).startswith('mps'			):
																snake_case	    =	torch.manual_seed(__a			)
										else:
																snake_case	    =	torch.Generator(device=__a			).manual_seed(__a			)

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

										return inputs

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	'cpu'  # ensure determinism for the device-dependent torch.Generator
										snake_case	    =	self.get_dummy_components()
										snake_case	    =	ConsistencyModelPipeline(**__a			)
										snake_case	    =	pipe.to(__a			)
										pipe.set_progress_bar_config(disable=__a			)

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

										snake_case	    =	image[0, -3:, -3:, -1]
										snake_case	    =	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 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	'cpu'  # ensure determinism for the device-dependent torch.Generator
										snake_case	    =	self.get_dummy_components(class_cond=__a			)
										snake_case	    =	ConsistencyModelPipeline(**__a			)
										snake_case	    =	pipe.to(__a			)
										pipe.set_progress_bar_config(disable=__a			)

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

										snake_case	    =	image[0, -3:, -3:, -1]
										snake_case	    =	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 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	'cpu'  # ensure determinism for the device-dependent torch.Generator
										snake_case	    =	self.get_dummy_components()
										snake_case	    =	ConsistencyModelPipeline(**__a			)
										snake_case	    =	pipe.to(__a			)
										pipe.set_progress_bar_config(disable=__a			)

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

										snake_case	    =	image[0, -3:, -3:, -1]
										snake_case	    =	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 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	'cpu'  # ensure determinism for the device-dependent torch.Generator
										snake_case	    =	self.get_dummy_components(class_cond=__a			)
										snake_case	    =	ConsistencyModelPipeline(**__a			)
										snake_case	    =	pipe.to(__a			)
										pipe.set_progress_bar_config(disable=__a			)

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

										snake_case	    =	image[0, -3:, -3:, -1]
										snake_case	    =	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 						snake_case       ( self			):






										"""simple docstring"""

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

				def 						snake_case       ( self						,			lowerCAmelCase=0						,			lowerCAmelCase=False						,			lowerCAmelCase="cpu"						,			lowerCAmelCase=torch.floataa						,			lowerCAmelCase=(1, 3, 64, 64)			):






										"""simple docstring"""

										snake_case	    =	torch.manual_seed(__a			)

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

										if get_fixed_latents:
																snake_case	    =	self.get_fixed_latents(seed=__a						,			device=__a						,			dtype=__a						,			shape=__a			)
																snake_case	    =	latents

										return inputs

				def 						snake_case       ( self						,			lowerCAmelCase=0						,			lowerCAmelCase="cpu"						,			lowerCAmelCase=torch.floataa						,			lowerCAmelCase=(1, 3, 64, 64)			):






										"""simple docstring"""

										if type(__a			) == str:
																snake_case	    =	torch.device(__a			)
										snake_case	    =	torch.Generator(device=__a			).manual_seed(__a			)
										snake_case	    =	randn_tensor(__a						,			generator=__a						,			device=__a						,			dtype=__a			)
										return latents

				def 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	UNetaDModel.from_pretrained('diffusers/consistency_models'						,			subfolder='diffusers_cd_imagenet64_l2'			)
										snake_case	    =	CMStochasticIterativeScheduler(
										    num_train_timesteps=40						,			sigma_min=0.0_02						,			sigma_max=80.0						,			)
										snake_case	    =	ConsistencyModelPipeline(unet=__a						,			scheduler=__a			)
										pipe.to(torch_device=__a			)
										pipe.set_progress_bar_config(disable=__a			)

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

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

										snake_case	    =	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 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	UNetaDModel.from_pretrained('diffusers/consistency_models'						,			subfolder='diffusers_cd_imagenet64_l2'			)
										snake_case	    =	CMStochasticIterativeScheduler(
										    num_train_timesteps=40						,			sigma_min=0.0_02						,			sigma_max=80.0						,			)
										snake_case	    =	ConsistencyModelPipeline(unet=__a						,			scheduler=__a			)
										pipe.to(torch_device=__a			)
										pipe.set_progress_bar_config(disable=__a			)

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

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

										snake_case	    =	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 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	UNetaDModel.from_pretrained('diffusers/consistency_models'						,			subfolder='diffusers_cd_imagenet64_l2'			)
										snake_case	    =	CMStochasticIterativeScheduler(
										    num_train_timesteps=40						,			sigma_min=0.0_02						,			sigma_max=80.0						,			)
										snake_case	    =	ConsistencyModelPipeline(unet=__a						,			scheduler=__a			)
										pipe.to(torch_device=__a						,			torch_dtype=torch.floataa			)
										pipe.set_progress_bar_config(disable=__a			)

										snake_case	    =	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			):
																snake_case	    =	pipe(**__a			).images
										assert image.shape == (1, 64, 64, 3)

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

										snake_case	    =	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 						snake_case       ( self			):






										"""simple docstring"""

										snake_case	    =	UNetaDModel.from_pretrained('diffusers/consistency_models'						,			subfolder='diffusers_cd_imagenet64_l2'			)
										snake_case	    =	CMStochasticIterativeScheduler(
										    num_train_timesteps=40						,			sigma_min=0.0_02						,			sigma_max=80.0						,			)
										snake_case	    =	ConsistencyModelPipeline(unet=__a						,			scheduler=__a			)
										pipe.to(torch_device=__a						,			torch_dtype=torch.floataa			)
										pipe.set_progress_bar_config(disable=__a			)

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

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

										snake_case	    =	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



'''simple docstring'''





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

import datasets


SCREAMING_SNAKE_CASE_:     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'


SCREAMING_SNAKE_CASE_:     Union[str, 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'


SCREAMING_SNAKE_CASE_:     List[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        __A      ( datasets.Metric	):



    def 					_lowercase						(self    :    Optional[Any] ):
         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 					_lowercase						(self    :    Tuple   ,			__a    :    Optional[int]   ,			__a    :    List[Any] ):
         UpperCAmelCase_           =   0.0
         for i, j in zip(__a   ,			__a ):
              n_correct += 1.0 if math_equivalence.is_equiv(__a   ,			__a ) else 0.0
         UpperCAmelCase_           =   n_correct / len(__a )
         return {
             "accuracy": accuracy,
         }




