




'''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  __lowercase    (	_lowercase							,				unittest.TestCase			):
   _a    						=   ConsistencyModelPipeline
   _a    						=   UNCONDITIONAL_IMAGE_GENERATION_PARAMS
   _a    						=   UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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







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







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







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase=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.002   ,				sigma_max=80.0   ,				)

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

     return components







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase=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 		UpperCamelCase__      (						self   )					->     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.3_572, 0.6_273, 0.4_031, 0.3_961, 0.4_321, 0.5_730, 0.5_266, 0.4_780, 0.5_004]   )

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







   def 		UpperCamelCase__      (						self   )					->     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.3_572, 0.6_273, 0.4_031, 0.3_961, 0.4_321, 0.5_730, 0.5_266, 0.4_780, 0.5_004]   )

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







   def 		UpperCamelCase__      (						self   )					->     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.5_004, 0.5_004, 0.4_994, 0.5_008, 0.4_976, 0.5_018, 0.4_990, 0.4_982, 0.4_987]   )

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







   def 		UpperCamelCase__      (						self   )					->     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.5_004, 0.5_004, 0.4_994, 0.5_008, 0.4_976, 0.5_018, 0.4_990, 0.4_982, 0.4_987]   )

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





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







   def 		UpperCamelCase__      (						self   )					->     str:
     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)   )					->     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 		UpperCamelCase__      (						self   ,				UpperCamelCase=0   ,				UpperCamelCase="cpu"   ,				UpperCamelCase=torch.floataa   ,				UpperCamelCase=(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 		UpperCamelCase__      (						self   )					->     List[str]:
     __a								=       UNetaDModel.from_pretrained('diffusers/consistency_models'   ,				subfolder='diffusers_cd_imagenet64_l2'   )
     __a								=       CMStochasticIterativeScheduler(
         num_train_timesteps=40   ,				sigma_min=0.002   ,				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.0_888, 0.0_881, 0.0_666, 0.0_479, 0.0_292, 0.0_195, 0.0_201, 0.0_163, 0.0_254]   )

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







   def 		UpperCamelCase__      (						self   )					->     List[str]:
     __a								=       UNetaDModel.from_pretrained('diffusers/consistency_models'   ,				subfolder='diffusers_cd_imagenet64_l2'   )
     __a								=       CMStochasticIterativeScheduler(
         num_train_timesteps=40   ,				sigma_min=0.002   ,				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.0_340, 0.0_152, 0.0_063, 0.0_267, 0.0_221, 0.0_107, 0.0_416, 0.0_186, 0.0_217]   )

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







   @require_torch_a
   def 		UpperCamelCase__      (						self   )					->     List[str]:
     __a								=       UNetaDModel.from_pretrained('diffusers/consistency_models'   ,				subfolder='diffusers_cd_imagenet64_l2'   )
     __a								=       CMStochasticIterativeScheduler(
         num_train_timesteps=40   ,				sigma_min=0.002   ,				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.1_875, 0.1_428, 0.1_289, 0.2_151, 0.2_092, 0.1_477, 0.1_877, 0.1_641, 0.1_353]   )

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







   @require_torch_a
   def 		UpperCamelCase__      (						self   )					->     List[Any]:
     __a								=       UNetaDModel.from_pretrained('diffusers/consistency_models'   ,				subfolder='diffusers_cd_imagenet64_l2'   )
     __a								=       CMStochasticIterativeScheduler(
         num_train_timesteps=40   ,				sigma_min=0.002   ,				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.1_663, 0.1_948, 0.2_275, 0.1_680, 0.1_204, 0.1_245, 0.1_858, 0.1_338, 0.2_095]   )

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





"""simple docstring"""


import os

# Precomputes a list of the 100 first triangular numbers
_lowercase				 =					[int(0.5 * n * (n + 1)) for n in range(1, 1_01)]






def    _snake_case					(		):
	A     = os.path.dirname(os.path.realpath(snake_case__  )  )
	A     = os.path.join(snake_case__						,   'words.txt'  )

	A     = ''
	with open(snake_case__  ) as f:
		A     = f.readline()

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


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