
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   UpperCAmelCase		(						A_       ,unittest.TestCase ):
   A__					:		Union[str, Any]						     =     ConsistencyModelPipeline
   A__					:		Union[str, Any]						     =     UNCONDITIONAL_IMAGE_GENERATION_PARAMS
   A__					:		Any						     =     UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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



   @property
   def 				_SCREAMING_SNAKE_CASE   (self		:					Dict )		->	Union[str, Any]:


          '''simple docstring'''


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



   @property
   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[Any] )		->	Union[str, Any]:


          '''simple docstring'''


          snake_case       :   List[Any]							=      UNetaDModel.from_pretrained(
              "diffusers/consistency-models-test"    ,						subfolder="test_unet_class_cond"    ,						)
          return unet



   def 				_SCREAMING_SNAKE_CASE   (self		:					List[str]    ,						snake_case__		:					Union[str, Any]=False )		->	List[Any]:


          '''simple docstring'''


          if class_cond:
                 snake_case       :   List[Any]							=      self.dummy_cond_unet
          else:
                 snake_case       :   Optional[int]							=      self.dummy_uncond_unet

          # Default to CM multistep sampler
          snake_case       :   Any							=      CMStochasticIterativeScheduler(
              num_train_timesteps=40    ,						sigma_min=0.002    ,						sigma_max=80.0    ,						)

          snake_case       :   List[Any]							=      {
              "unet": unet,
              "scheduler": scheduler,
          }

          return components



   def 				_SCREAMING_SNAKE_CASE   (self		:					str    ,						snake_case__		:					Any    ,						snake_case__		:					Optional[Any]=0 )		->	Optional[Any]:


          '''simple docstring'''


          if str(snake_case__ ).startswith("mps" ):
                 snake_case       :   List[str]							=      torch.manual_seed(snake_case__ )
          else:
                 snake_case       :   Any							=      torch.Generator(device=snake_case__ ).manual_seed(snake_case__ )

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

          return inputs



   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[int] )		->	str:


          '''simple docstring'''


          snake_case       :   Optional[int]							=      "cpu"  # ensure determinism for the device-dependent torch.Generator
          snake_case       :   Dict							=      self.get_dummy_components()
          snake_case       :   Optional[Any]							=      ConsistencyModelPipeline(**snake_case__ )
          snake_case       :   str							=      pipe.to(snake_case__ )
          pipe.set_progress_bar_config(disable=snake_case__ )

          snake_case       :   Union[str, Any]							=      self.get_dummy_inputs(snake_case__ )
          snake_case       :   Optional[int]							=      pipe(**snake_case__ ).images
          assert image.shape == (1, 32, 32, 3)

          snake_case       :   Union[str, Any]							=      image[0, -3:, -3:, -1]
          snake_case       :   Dict							=      np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004] )

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



   def 				_SCREAMING_SNAKE_CASE   (self		:					int )		->	Any:


          '''simple docstring'''


          snake_case       :   List[Any]							=      "cpu"  # ensure determinism for the device-dependent torch.Generator
          snake_case       :   int							=      self.get_dummy_components(class_cond=snake_case__ )
          snake_case       :   Any							=      ConsistencyModelPipeline(**snake_case__ )
          snake_case       :   List[Any]							=      pipe.to(snake_case__ )
          pipe.set_progress_bar_config(disable=snake_case__ )

          snake_case       :   Union[str, Any]							=      self.get_dummy_inputs(snake_case__ )
          snake_case       :   Dict							=      0
          snake_case       :   List[Any]							=      pipe(**snake_case__ ).images
          assert image.shape == (1, 32, 32, 3)

          snake_case       :   Optional[Any]							=      image[0, -3:, -3:, -1]
          snake_case       :   Optional[Any]							=      np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004] )

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



   def 				_SCREAMING_SNAKE_CASE   (self		:					Tuple )		->	Optional[int]:


          '''simple docstring'''


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

          snake_case       :   Dict							=      self.get_dummy_inputs(snake_case__ )
          snake_case       :   List[Any]							=      1
          snake_case       :   Dict							=      None
          snake_case       :   Tuple							=      pipe(**snake_case__ ).images
          assert image.shape == (1, 32, 32, 3)

          snake_case       :   Any							=      image[0, -3:, -3:, -1]
          snake_case       :   List[str]							=      np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987] )

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



   def 				_SCREAMING_SNAKE_CASE   (self		:					List[Any] )		->	Any:


          '''simple docstring'''


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

          snake_case       :   List[str]							=      self.get_dummy_inputs(snake_case__ )
          snake_case       :   Any							=      1
          snake_case       :   Optional[int]							=      None
          snake_case       :   List[Any]							=      0
          snake_case       :   int							=      pipe(**snake_case__ ).images
          assert image.shape == (1, 32, 32, 3)

          snake_case       :   Union[str, Any]							=      image[0, -3:, -3:, -1]
          snake_case       :   Dict							=      np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987] )

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






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



   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[Any] )		->	Tuple:


          '''simple docstring'''


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



   def 				_SCREAMING_SNAKE_CASE   (self		:					Any    ,						snake_case__		:					Union[str, Any]=0    ,						snake_case__		:					Optional[Any]=False    ,						snake_case__		:					str="cpu"    ,						snake_case__		:					Optional[int]=torch.floataa    ,						snake_case__		:					Optional[Any]=(1, 3, 64, 64) )		->	Tuple:


          '''simple docstring'''


          snake_case       :   List[str]							=      torch.manual_seed(snake_case__ )

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

          if get_fixed_latents:
                 snake_case       :   Tuple							=      self.get_fixed_latents(seed=snake_case__    ,						device=snake_case__    ,						dtype=snake_case__    ,						shape=snake_case__ )
                 snake_case       :   int							=      latents

          return inputs



   def 				_SCREAMING_SNAKE_CASE   (self		:					Any    ,						snake_case__		:					List[Any]=0    ,						snake_case__		:					Union[str, Any]="cpu"    ,						snake_case__		:					int=torch.floataa    ,						snake_case__		:					Any=(1, 3, 64, 64) )		->	int:


          '''simple docstring'''


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



   def 				_SCREAMING_SNAKE_CASE   (self		:					Any )		->	Union[str, Any]:


          '''simple docstring'''


          snake_case       :   Optional[Any]							=      UNetaDModel.from_pretrained("diffusers/consistency_models"    ,						subfolder="diffusers_cd_imagenet64_l2" )
          snake_case       :   Tuple							=      CMStochasticIterativeScheduler(
              num_train_timesteps=40    ,						sigma_min=0.002    ,						sigma_max=80.0    ,						)
          snake_case       :   List[Any]							=      ConsistencyModelPipeline(unet=snake_case__    ,						scheduler=snake_case__ )
          pipe.to(torch_device=snake_case__ )
          pipe.set_progress_bar_config(disable=snake_case__ )

          snake_case       :   str							=      self.get_inputs()
          snake_case       :   Optional[int]							=      pipe(**snake_case__ ).images
          assert image.shape == (1, 64, 64, 3)

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

          snake_case       :   List[str]							=      np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254] )

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



   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[int] )		->	Union[str, Any]:


          '''simple docstring'''


          snake_case       :   Dict							=      UNetaDModel.from_pretrained("diffusers/consistency_models"    ,						subfolder="diffusers_cd_imagenet64_l2" )
          snake_case       :   Tuple							=      CMStochasticIterativeScheduler(
              num_train_timesteps=40    ,						sigma_min=0.002    ,						sigma_max=80.0    ,						)
          snake_case       :   Dict							=      ConsistencyModelPipeline(unet=snake_case__    ,						scheduler=snake_case__ )
          pipe.to(torch_device=snake_case__ )
          pipe.set_progress_bar_config(disable=snake_case__ )

          snake_case       :   int							=      self.get_inputs()
          snake_case       :   Union[str, Any]							=      1
          snake_case       :   Optional[int]							=      None
          snake_case       :   int							=      pipe(**snake_case__ ).images
          assert image.shape == (1, 64, 64, 3)

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

          snake_case       :   Union[str, Any]							=      np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217] )

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



   @require_torch_a
   def 				_SCREAMING_SNAKE_CASE   (self		:					Union[str, Any] )		->	Union[str, Any]:


          '''simple docstring'''


          snake_case       :   str							=      UNetaDModel.from_pretrained("diffusers/consistency_models"    ,						subfolder="diffusers_cd_imagenet64_l2" )
          snake_case       :   Optional[int]							=      CMStochasticIterativeScheduler(
              num_train_timesteps=40    ,						sigma_min=0.002    ,						sigma_max=80.0    ,						)
          snake_case       :   Optional[int]							=      ConsistencyModelPipeline(unet=snake_case__    ,						scheduler=snake_case__ )
          pipe.to(torch_device=snake_case__    ,						torch_dtype=torch.floataa )
          pipe.set_progress_bar_config(disable=snake_case__ )

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

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

          snake_case       :   Optional[Any]							=      np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353] )

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



   @require_torch_a
   def 				_SCREAMING_SNAKE_CASE   (self		:					str )		->	List[Any]:


          '''simple docstring'''


          snake_case       :   Optional[int]							=      UNetaDModel.from_pretrained("diffusers/consistency_models"    ,						subfolder="diffusers_cd_imagenet64_l2" )
          snake_case       :   str							=      CMStochasticIterativeScheduler(
              num_train_timesteps=40    ,						sigma_min=0.002    ,						sigma_max=80.0    ,						)
          snake_case       :   str							=      ConsistencyModelPipeline(unet=snake_case__    ,						scheduler=snake_case__ )
          pipe.to(torch_device=snake_case__    ,						torch_dtype=torch.floataa )
          pipe.set_progress_bar_config(disable=snake_case__ )

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

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

          snake_case       :   List[Any]							=      np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095] )

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







__lowerCamelCase							=       {
    "joule": 1.0,
    "kilojoule": 10_00,
    "megajoule": 1_00_00_00,
    "gigajoule": 10_00_00_00_00,
    "wattsecond": 1.0,
    "watthour": 36_00,
    "kilowatthour": 3_60_00_00,
    "newtonmeter": 1.0,
    "calorie_nutr": 41_86.8,
    "kilocalorie_nutr": 4_18_68_00.00,
    "electronvolt": 1.602_176_634e-19,
    "britishthermalunit_it": 10_55.0_55_85,
    "footpound": 1.35_5818,
}


def 			UpperCamelCase					(    __lowerCamelCase     : str				,					__lowerCamelCase     : str				,					__lowerCamelCase     : float    ):
       if to_type not in ENERGY_CONVERSION or from_type not in ENERGY_CONVERSION:
              snake_case       :   List[Any]							=      (
                  f"""Incorrect 'from_type' or 'to_type' value: {from_type!r}, {to_type!r}\n"""
                  f"""Valid values are: {', '.join(__lowerCamelCase    )}"""
              )
              raise ValueError(__lowerCamelCase    )
       return value * ENERGY_CONVERSION[from_type] / ENERGY_CONVERSION[to_type]


if __name__ == "__main__":
   import doctest

   doctest.testmod()





