
"""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 re
import string

import numpy as np

import datasets


__UpperCamelCase     		=							'\nReturns the rate at which the input predicted strings exactly match their references, ignoring any strings input as part of the regexes_to_ignore list.\n'

__UpperCamelCase     		=							'\nArgs:\n    predictions: List of predicted texts.\n    references: List of reference texts.\n    regexes_to_ignore: List, defaults to None. Regex expressions of characters to\n        ignore when calculating the exact matches. Note: these regexes are removed\n        from the input data before the changes based on the options below (e.g. ignore_case,\n        ignore_punctuation, ignore_numbers) are applied.\n    ignore_case: Boolean, defaults to False. If true, turns everything\n        to lowercase so that capitalization differences are ignored.\n    ignore_punctuation: Boolean, defaults to False. If true, removes all punctuation before\n        comparing predictions and references.\n    ignore_numbers: Boolean, defaults to False. If true, removes all punctuation before\n        comparing predictions and references.\nReturns:\n    exact_match: Dictionary containing exact_match rate. Possible values are between 0.0 and 100.0, inclusive.\nExamples:\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds)\n    >>> print(round(results["exact_match"], 1))\n    25.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell"], ignore_case=True, ignore_punctuation=True)\n    >>> print(round(results["exact_match"], 1))\n    50.0\n\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell", "YELL"], ignore_case=True, ignore_punctuation=True)\n    >>> print(round(results["exact_match"], 1))\n    75.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell", "YELL"], ignore_case=True, ignore_punctuation=True, ignore_numbers=True)\n    >>> print(round(results["exact_match"], 1))\n    100.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["The cat sat on the mat.", "Theaters are great.", "It\'s like comparing oranges and apples."]\n    >>> preds = ["The cat sat on the mat?", "Theaters are great.", "It\'s like comparing apples and oranges."]\n    >>> results = exact_match.compute(references=refs, predictions=preds)\n    >>> print(round(results["exact_match"], 1))\n    33.3\n\n'

__UpperCamelCase     		=							'\n'







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






			"""simple docstring"""






			def 					snake_case__					(				self  ):

						'''simple docstring'''


						return datasets.MetricInfo(
						    description=_DESCRIPTION   ,    citation=_CITATION   ,    inputs_description=_KWARGS_DESCRIPTION   ,    features=datasets.Features(
						        {
						            "predictions": datasets.Value("string"   ,    id="sequence"  ),
						            "references": datasets.Value("string"   ,    id="sequence"  ),
						        }  )   ,    reference_urls=[]   ,    )







			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case=None   ,    snake_case=False   ,    snake_case=False   ,    snake_case=False   ,    ):

						'''simple docstring'''


						if regexes_to_ignore is not None:
									for s in regexes_to_ignore:
												UpperCamelCase__       =							np.array([re.sub(snake_case   ,    ""   ,    snake_case  ) for x in predictions]  )
												UpperCamelCase__       =							np.array([re.sub(snake_case   ,    ""   ,    snake_case  ) for x in references]  )
						else:
									UpperCamelCase__       =							np.asarray(snake_case  )
									UpperCamelCase__       =							np.asarray(snake_case  )

						if ignore_case:
									UpperCamelCase__       =							np.char.lower(snake_case  )
									UpperCamelCase__       =							np.char.lower(snake_case  )

						if ignore_punctuation:
									UpperCamelCase__       =							string.punctuation.maketrans(""   ,    ""   ,    string.punctuation  )
									UpperCamelCase__       =							np.char.translate(snake_case   ,    table=snake_case  )
									UpperCamelCase__       =							np.char.translate(snake_case   ,    table=snake_case  )

						if ignore_numbers:
									UpperCamelCase__       =							string.digits.maketrans(""   ,    ""   ,    string.digits  )
									UpperCamelCase__       =							np.char.translate(snake_case   ,    table=snake_case  )
									UpperCamelCase__       =							np.char.translate(snake_case   ,    table=snake_case  )

						UpperCamelCase__       =							predictions == references

						return {"exact_match": np.mean(snake_case  ) * 100}


