





"""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   (     a  ,   unittest.TestCase       ):
       """simple docstring"""




       __magic_name__     :Optional[Any]		  =					ConsistencyModelPipeline
       __magic_name__     :Any		  =					UNCONDITIONAL_IMAGE_GENERATION_PARAMS
       __magic_name__     :Dict		  =					UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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






       @property
       def      snake_case						(    self      ):



         '''simple docstring'''

         lowerCAmelCase__			:List[Any]											=							UNetaDModel.from_pretrained(
             'diffusers/consistency-models-test'    , subfolder='test_unet'    , )
         return unet






       @property
       def      snake_case						(    self      ):



         '''simple docstring'''

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






       def      snake_case						(    self    , __UpperCAmelCase=False      ):



         '''simple docstring'''

         if class_cond:
           lowerCAmelCase__			:Dict											=							self.dummy_cond_unet
         else:
           lowerCAmelCase__			:List[str]											=							self.dummy_uncond_unet

         # Default to CM multistep sampler
         lowerCAmelCase__			:Any											=							CMStochasticIterativeScheduler(
             num_train_timesteps=4_0    , sigma_min=0.0_02    , sigma_max=80.0    , )

         lowerCAmelCase__			:List[Any]											=							{
             'unet': unet,
             'scheduler': scheduler,
         }

         return components






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



         '''simple docstring'''

         if str(__UpperCAmelCase      ).startswith('mps'      ):
           lowerCAmelCase__			:Dict											=							torch.manual_seed(__UpperCAmelCase      )
         else:
           lowerCAmelCase__			:Any											=							torch.Generator(device=__UpperCAmelCase      ).manual_seed(__UpperCAmelCase      )

         lowerCAmelCase__			:List[Any]											=							{
             'batch_size': 1,
             'num_inference_steps': None,
             'timesteps': [2_2, 0],
             'generator': generator,
             'output_type': 'np',
         }

         return inputs






       def      snake_case						(    self      ):



         '''simple docstring'''

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

         lowerCAmelCase__			:Optional[int]											=							self.get_dummy_inputs(__UpperCAmelCase      )
         lowerCAmelCase__			:Tuple											=							pipe(**__UpperCAmelCase      ).images
         assert image.shape == (1, 3_2, 3_2, 3)

         lowerCAmelCase__			:Tuple											=							image[0, -3:, -3:, -1]
         lowerCAmelCase__			:Tuple											=							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'''

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

         lowerCAmelCase__			:Tuple											=							self.get_dummy_inputs(__UpperCAmelCase      )
         lowerCAmelCase__			:str											=							0
         lowerCAmelCase__			:Union[str, Any]											=							pipe(**__UpperCAmelCase      ).images
         assert image.shape == (1, 3_2, 3_2, 3)

         lowerCAmelCase__			:Any											=							image[0, -3:, -3:, -1]
         lowerCAmelCase__			:Dict											=							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'''

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

         lowerCAmelCase__			:List[str]											=							self.get_dummy_inputs(__UpperCAmelCase      )
         lowerCAmelCase__			:Union[str, Any]											=							1
         lowerCAmelCase__			:Optional[Any]											=							None
         lowerCAmelCase__			:Any											=							pipe(**__UpperCAmelCase      ).images
         assert image.shape == (1, 3_2, 3_2, 3)

         lowerCAmelCase__			:Tuple											=							image[0, -3:, -3:, -1]
         lowerCAmelCase__			:List[Any]											=							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'''

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

         lowerCAmelCase__			:Any											=							self.get_dummy_inputs(__UpperCAmelCase      )
         lowerCAmelCase__			:List[str]											=							1
         lowerCAmelCase__			:Dict											=							None
         lowerCAmelCase__			:int											=							0
         lowerCAmelCase__			:List[Any]											=							pipe(**__UpperCAmelCase      ).images
         assert image.shape == (1, 3_2, 3_2, 3)

         lowerCAmelCase__			:int											=							image[0, -3:, -3:, -1]
         lowerCAmelCase__			:Union[str, Any]											=							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    , __UpperCAmelCase=0    , __UpperCAmelCase=False    , __UpperCAmelCase="cpu"    , __UpperCAmelCase=torch.floataa    , __UpperCAmelCase=(1, 3, 6_4, 6_4)      ):



         '''simple docstring'''

         lowerCAmelCase__			:str											=							torch.manual_seed(__UpperCAmelCase      )

         lowerCAmelCase__			:Tuple											=							{
             'num_inference_steps': None,
             'timesteps': [2_2, 0],
             'class_labels': 0,
             'generator': generator,
             'output_type': 'np',
         }

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

         return inputs






       def      snake_case						(    self    , __UpperCAmelCase=0    , __UpperCAmelCase="cpu"    , __UpperCAmelCase=torch.floataa    , __UpperCAmelCase=(1, 3, 6_4, 6_4)      ):



         '''simple docstring'''

         if type(__UpperCAmelCase      ) == str:
           lowerCAmelCase__			:Tuple											=							torch.device(__UpperCAmelCase      )
         lowerCAmelCase__			:List[str]											=							torch.Generator(device=__UpperCAmelCase      ).manual_seed(__UpperCAmelCase      )
         lowerCAmelCase__			:Optional[Any]											=							randn_tensor(__UpperCAmelCase    , generator=__UpperCAmelCase    , device=__UpperCAmelCase    , dtype=__UpperCAmelCase      )
         return latents






       def      snake_case						(    self      ):



         '''simple docstring'''

         lowerCAmelCase__			:str											=							UNetaDModel.from_pretrained('diffusers/consistency_models'    , subfolder='diffusers_cd_imagenet64_l2'      )
         lowerCAmelCase__			:List[str]											=							CMStochasticIterativeScheduler(
             num_train_timesteps=4_0    , sigma_min=0.0_02    , sigma_max=80.0    , )
         lowerCAmelCase__			:List[Any]											=							ConsistencyModelPipeline(unet=__UpperCAmelCase    , scheduler=__UpperCAmelCase      )
         pipe.to(torch_device=__UpperCAmelCase      )
         pipe.set_progress_bar_config(disable=__UpperCAmelCase      )

         lowerCAmelCase__			:int											=							self.get_inputs()
         lowerCAmelCase__			:int											=							pipe(**__UpperCAmelCase      ).images
         assert image.shape == (1, 6_4, 6_4, 3)

         lowerCAmelCase__			:Any											=							image[0, -3:, -3:, -1]

         lowerCAmelCase__			:Tuple											=							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'''

         lowerCAmelCase__			:Dict											=							UNetaDModel.from_pretrained('diffusers/consistency_models'    , subfolder='diffusers_cd_imagenet64_l2'      )
         lowerCAmelCase__			:str											=							CMStochasticIterativeScheduler(
             num_train_timesteps=4_0    , sigma_min=0.0_02    , sigma_max=80.0    , )
         lowerCAmelCase__			:Optional[Any]											=							ConsistencyModelPipeline(unet=__UpperCAmelCase    , scheduler=__UpperCAmelCase      )
         pipe.to(torch_device=__UpperCAmelCase      )
         pipe.set_progress_bar_config(disable=__UpperCAmelCase      )

         lowerCAmelCase__			:List[str]											=							self.get_inputs()
         lowerCAmelCase__			:Tuple											=							1
         lowerCAmelCase__			:int											=							None
         lowerCAmelCase__			:Optional[int]											=							pipe(**__UpperCAmelCase      ).images
         assert image.shape == (1, 6_4, 6_4, 3)

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

         lowerCAmelCase__			:Union[str, Any]											=							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'''

         lowerCAmelCase__			:Tuple											=							UNetaDModel.from_pretrained('diffusers/consistency_models'    , subfolder='diffusers_cd_imagenet64_l2'      )
         lowerCAmelCase__			:Dict											=							CMStochasticIterativeScheduler(
             num_train_timesteps=4_0    , sigma_min=0.0_02    , sigma_max=80.0    , )
         lowerCAmelCase__			:Tuple											=							ConsistencyModelPipeline(unet=__UpperCAmelCase    , scheduler=__UpperCAmelCase      )
         pipe.to(torch_device=__UpperCAmelCase    , torch_dtype=torch.floataa      )
         pipe.set_progress_bar_config(disable=__UpperCAmelCase      )

         lowerCAmelCase__			:Dict											=							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      ):
           lowerCAmelCase__			:List[str]											=							pipe(**__UpperCAmelCase      ).images
         assert image.shape == (1, 6_4, 6_4, 3)

         lowerCAmelCase__			:Any											=							image[0, -3:, -3:, -1]

         lowerCAmelCase__			:int											=							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'''

         lowerCAmelCase__			:Any											=							UNetaDModel.from_pretrained('diffusers/consistency_models'    , subfolder='diffusers_cd_imagenet64_l2'      )
         lowerCAmelCase__			:Dict											=							CMStochasticIterativeScheduler(
             num_train_timesteps=4_0    , sigma_min=0.0_02    , sigma_max=80.0    , )
         lowerCAmelCase__			:Optional[Any]											=							ConsistencyModelPipeline(unet=__UpperCAmelCase    , scheduler=__UpperCAmelCase      )
         pipe.to(torch_device=__UpperCAmelCase    , torch_dtype=torch.floataa      )
         pipe.set_progress_bar_config(disable=__UpperCAmelCase      )

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

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

         lowerCAmelCase__			:List[Any]											=							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
def  __A   (_SCREAMING_SNAKE_CASE			)					->int:





  """simple docstring"""


  if not isinstance(_SCREAMING_SNAKE_CASE     ,   _SCREAMING_SNAKE_CASE			):
    lowerCAmelCase__			:Dict											=							F"Input value of [number={number}] must be an integer"
    raise TypeError(_SCREAMING_SNAKE_CASE			)

  if number < 1:
    lowerCAmelCase__			:Dict											=							F"Input value of [number={number}] must be > 0"
    raise ValueError(_SCREAMING_SNAKE_CASE			)
  elif number == 1:
    return 3
  elif number == 2:
    return 5
  else:
    lowerCAmelCase__			:Union[str, Any]											=							int(math.log(number // 3     ,   2			)			) + 2

    lowerCAmelCase__			:Optional[Any]											=							[3, 5]
    lowerCAmelCase__			:Optional[Any]											=							2
    lowerCAmelCase__			:List[str]											=							3
    for block in range(1     ,   _SCREAMING_SNAKE_CASE			):
      for _ in range(_SCREAMING_SNAKE_CASE			):
        proth_list.append(2 ** (block + 1) + proth_list[proth_index - 1]			)
        proth_index += 1
      increment *= 2

  return proth_list[number - 1]


if __name__ == "__main__":
  import doctest

  doctest.testmod()

  for number in range(11):
    __A											=		0
    try:
      __A											=		proth(number)
    except ValueError:
      print(F'''ValueError: there is no {number}th Proth number''')
      continue

    print(F'''The {number}th Proth number: {value}''')
