
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    a_ (			_a					,			unittest.TestCase    ):
			a	:			int					   =							ConsistencyModelPipeline
			a	:			Dict					   =							UNCONDITIONAL_IMAGE_GENERATION_PARAMS
			a	:			List[str]					   =							UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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






			@property
			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					UNetaDModel.from_pretrained(
									    """diffusers/consistency-models-test"""				,			subfolder="""test_unet"""				,			)
									return unet






			@property
			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					UNetaDModel.from_pretrained(
									    """diffusers/consistency-models-test"""				,			subfolder="""test_unet_class_cond"""				,			)
									return unet






			def    UpperCamelCase_				(      self				,			__UpperCamelCase=False  ):
									if class_cond:
															_lowercase					       =					self.dummy_cond_unet
									else:
															_lowercase					       =					self.dummy_uncond_unet

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

									_lowercase					       =					{
									    """unet""": unet,
									    """scheduler""": scheduler,
									}

									return components






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase=0  ):
									if str(__UpperCamelCase  ).startswith("""mps"""  ):
															_lowercase					       =					torch.manual_seed(__UpperCamelCase  )
									else:
															_lowercase					       =					torch.Generator(device=__UpperCamelCase  ).manual_seed(__UpperCamelCase  )

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

									return inputs






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
									_lowercase					       =					self.get_dummy_components()
									_lowercase					       =					ConsistencyModelPipeline(**__UpperCamelCase  )
									_lowercase					       =					pipe.to(__UpperCamelCase  )
									pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

									_lowercase					       =					image[0, -3:, -3:, -1]
									_lowercase					       =					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    UpperCamelCase_				(      self  ):
									_lowercase					       =					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
									_lowercase					       =					self.get_dummy_components(class_cond=__UpperCamelCase  )
									_lowercase					       =					ConsistencyModelPipeline(**__UpperCamelCase  )
									_lowercase					       =					pipe.to(__UpperCamelCase  )
									pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

									_lowercase					       =					image[0, -3:, -3:, -1]
									_lowercase					       =					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    UpperCamelCase_				(      self  ):
									_lowercase					       =					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
									_lowercase					       =					self.get_dummy_components()
									_lowercase					       =					ConsistencyModelPipeline(**__UpperCamelCase  )
									_lowercase					       =					pipe.to(__UpperCamelCase  )
									pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

									_lowercase					       =					image[0, -3:, -3:, -1]
									_lowercase					       =					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    UpperCamelCase_				(      self  ):
									_lowercase					       =					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
									_lowercase					       =					self.get_dummy_components(class_cond=__UpperCamelCase  )
									_lowercase					       =					ConsistencyModelPipeline(**__UpperCamelCase  )
									_lowercase					       =					pipe.to(__UpperCamelCase  )
									pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

									_lowercase					       =					image[0, -3:, -3:, -1]
									_lowercase					       =					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    a_ (			unittest.TestCase    ):






			def    UpperCamelCase_				(      self  ):
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()






			def    UpperCamelCase_				(      self				,			__UpperCamelCase=0				,			__UpperCamelCase=False				,			__UpperCamelCase="cpu"				,			__UpperCamelCase=torch.floataa				,			__UpperCamelCase=(1, 3, 64, 64)  ):
									_lowercase					       =					torch.manual_seed(__UpperCamelCase  )

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

									if get_fixed_latents:
															_lowercase					       =					self.get_fixed_latents(seed=__UpperCamelCase				,			device=__UpperCamelCase				,			dtype=__UpperCamelCase				,			shape=__UpperCamelCase  )
															_lowercase					       =					latents

									return inputs






			def    UpperCamelCase_				(      self				,			__UpperCamelCase=0				,			__UpperCamelCase="cpu"				,			__UpperCamelCase=torch.floataa				,			__UpperCamelCase=(1, 3, 64, 64)  ):
									if type(__UpperCamelCase  ) == str:
															_lowercase					       =					torch.device(__UpperCamelCase  )
									_lowercase					       =					torch.Generator(device=__UpperCamelCase  ).manual_seed(__UpperCamelCase  )
									_lowercase					       =					randn_tensor(__UpperCamelCase				,			generator=__UpperCamelCase				,			device=__UpperCamelCase				,			dtype=__UpperCamelCase  )
									return latents






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					UNetaDModel.from_pretrained("""diffusers/consistency_models"""				,			subfolder="""diffusers_cd_imagenet64_l2"""  )
									_lowercase					       =					CMStochasticIterativeScheduler(
									    num_train_timesteps=40				,			sigma_min=0.0_02				,			sigma_max=80.0				,			)
									_lowercase					       =					ConsistencyModelPipeline(unet=__UpperCamelCase				,			scheduler=__UpperCamelCase  )
									pipe.to(torch_device=__UpperCamelCase  )
									pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

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

									_lowercase					       =					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    UpperCamelCase_				(      self  ):
									_lowercase					       =					UNetaDModel.from_pretrained("""diffusers/consistency_models"""				,			subfolder="""diffusers_cd_imagenet64_l2"""  )
									_lowercase					       =					CMStochasticIterativeScheduler(
									    num_train_timesteps=40				,			sigma_min=0.0_02				,			sigma_max=80.0				,			)
									_lowercase					       =					ConsistencyModelPipeline(unet=__UpperCamelCase				,			scheduler=__UpperCamelCase  )
									pipe.to(torch_device=__UpperCamelCase  )
									pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

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

									_lowercase					       =					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    UpperCamelCase_				(      self  ):
									_lowercase					       =					UNetaDModel.from_pretrained("""diffusers/consistency_models"""				,			subfolder="""diffusers_cd_imagenet64_l2"""  )
									_lowercase					       =					CMStochasticIterativeScheduler(
									    num_train_timesteps=40				,			sigma_min=0.0_02				,			sigma_max=80.0				,			)
									_lowercase					       =					ConsistencyModelPipeline(unet=__UpperCamelCase				,			scheduler=__UpperCamelCase  )
									pipe.to(torch_device=__UpperCamelCase				,			torch_dtype=torch.floataa  )
									pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

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

									_lowercase					       =					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    UpperCamelCase_				(      self  ):
									_lowercase					       =					UNetaDModel.from_pretrained("""diffusers/consistency_models"""				,			subfolder="""diffusers_cd_imagenet64_l2"""  )
									_lowercase					       =					CMStochasticIterativeScheduler(
									    num_train_timesteps=40				,			sigma_min=0.0_02				,			sigma_max=80.0				,			)
									_lowercase					       =					ConsistencyModelPipeline(unet=__UpperCamelCase				,			scheduler=__UpperCamelCase  )
									pipe.to(torch_device=__UpperCamelCase				,			torch_dtype=torch.floataa  )
									pipe.set_progress_bar_config(disable=__UpperCamelCase  )

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

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

									_lowercase					       =					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 os

# Precomputes a list of the 100 first triangular numbers
A		: List[Any]								=     [int(0.5 * n * (n + 1)) for n in range(1, 101)]




def 	UpperCamelCase__  (     )  ->     str:
						_lowercase					       =					os.path.dirname(os.path.realpath(SCREAMING_SNAKE_CASE_       )       )
						_lowercase					       =					os.path.join(SCREAMING_SNAKE_CASE_		,   """words.txt"""       )

						_lowercase					       =					""""""
						with open(SCREAMING_SNAKE_CASE_       ) as f:
												_lowercase					       =					f.readline()

						_lowercase					       =					[word.strip("""\""""       ) for word in words.strip("""\r\n"""       ).split(""","""       )]
						_lowercase					       =					[
						    word
						    for word in [sum(ord(SCREAMING_SNAKE_CASE_       ) - 64 for x in word       ) for word in words]
						    if word in TRIANGULAR_NUMBERS
						]
						return len(SCREAMING_SNAKE_CASE_       )


if __name__ == "__main__":
		print(solution())