

'''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__     (					lowerCamelCase__				,			unittest.TestCase						):




     """simple docstring"""





     __UpperCamelCase									= ConsistencyModelPipeline
     __UpperCamelCase									= UNCONDITIONAL_IMAGE_GENERATION_PARAMS
     __UpperCamelCase									= UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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







     @property
     def 			__lowerCAmelCase	(		self						:	Dict	)       ->			str:
          '''simple docstring'''




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







     @property
     def 			__lowerCAmelCase	(		self						:	Any	)       ->			Optional[int]:
          '''simple docstring'''




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







     def 			__lowerCAmelCase	(		self						:	Any      ,						A__						:	str=False	)       ->			Optional[Any]:
          '''simple docstring'''




          if class_cond:
               a__   :      List[str]     =						self.dummy_cond_unet
          else:
               a__   :      Any     =						self.dummy_uncond_unet

          # Default to CM multistep sampler
          a__   :      int     =						CMStochasticIterativeScheduler(
              num_train_timesteps=4_0      ,						sigma_min=0.002      ,						sigma_max=8_0.0      ,						)

          a__   :      Dict     =						{
              '''unet''': unet,
              '''scheduler''': scheduler,
          }

          return components







     def 			__lowerCAmelCase	(		self						:	List[Any]      ,						A__						:	Tuple      ,						A__						:	Dict=0	)       ->			List[Any]:
          '''simple docstring'''




          if str(SCREAMING_SNAKE_CASE__	).startswith('''mps'''	):
               a__   :      Optional[Any]     =						torch.manual_seed(SCREAMING_SNAKE_CASE__	)
          else:
               a__   :      Optional[int]     =						torch.Generator(device=SCREAMING_SNAKE_CASE__	).manual_seed(SCREAMING_SNAKE_CASE__	)

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

          return inputs







     def 			__lowerCAmelCase	(		self						:	Optional[Any]	)       ->			Tuple:
          '''simple docstring'''




          a__   :      Tuple     =						'''cpu'''  # ensure determinism for the device-dependent torch.Generator
          a__   :      Dict     =						self.get_dummy_components()
          a__   :      Dict     =						ConsistencyModelPipeline(**SCREAMING_SNAKE_CASE__	)
          a__   :      List[str]     =						pipe.to(SCREAMING_SNAKE_CASE__	)
          pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__	)

          a__   :      Optional[Any]     =						self.get_dummy_inputs(SCREAMING_SNAKE_CASE__	)
          a__   :      Optional[Any]     =						pipe(**SCREAMING_SNAKE_CASE__	).images
          assert image.shape == (1, 3_2, 3_2, 3)

          a__   :      Tuple     =						image[0, -3:, -3:, -1]
          a__   :      List[Any]     =						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 			__lowerCAmelCase	(		self						:	Dict	)       ->			List[Any]:
          '''simple docstring'''




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

          a__   :      Optional[int]     =						self.get_dummy_inputs(SCREAMING_SNAKE_CASE__	)
          a__   :      int     =						0
          a__   :      Optional[int]     =						pipe(**SCREAMING_SNAKE_CASE__	).images
          assert image.shape == (1, 3_2, 3_2, 3)

          a__   :      Optional[Any]     =						image[0, -3:, -3:, -1]
          a__   :      Tuple     =						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 			__lowerCAmelCase	(		self						:	int	)       ->			Tuple:
          '''simple docstring'''




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

          a__   :      str     =						self.get_dummy_inputs(SCREAMING_SNAKE_CASE__	)
          a__   :      Union[str, Any]     =						1
          a__   :      Union[str, Any]     =						None
          a__   :      str     =						pipe(**SCREAMING_SNAKE_CASE__	).images
          assert image.shape == (1, 3_2, 3_2, 3)

          a__   :      Tuple     =						image[0, -3:, -3:, -1]
          a__   :      Dict     =						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 			__lowerCAmelCase	(		self						:	List[Any]	)       ->			Any:
          '''simple docstring'''




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

          a__   :      List[Any]     =						self.get_dummy_inputs(SCREAMING_SNAKE_CASE__	)
          a__   :      Optional[Any]     =						1
          a__   :      Optional[int]     =						None
          a__   :      Optional[Any]     =						0
          a__   :      Optional[int]     =						pipe(**SCREAMING_SNAKE_CASE__	).images
          assert image.shape == (1, 3_2, 3_2, 3)

          a__   :      List[str]     =						image[0, -3:, -3:, -1]
          a__   :      Any     =						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 			lowerCAmelCase__     (					unittest.TestCase						):




     """simple docstring"""





     def 			__lowerCAmelCase	(		self						:	int	)       ->			Optional[Any]:
          '''simple docstring'''




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







     def 			__lowerCAmelCase	(		self						:	List[str]      ,						A__						:	Optional[Any]=0      ,						A__						:	Tuple=False      ,						A__						:	Any="cpu"      ,						A__						:	Any=torch.floataa      ,						A__						:	List[Any]=(1, 3, 6_4, 6_4)	)       ->			Union[str, Any]:
          '''simple docstring'''




          a__   :      List[str]     =						torch.manual_seed(SCREAMING_SNAKE_CASE__	)

          a__   :      Optional[int]     =						{
              '''num_inference_steps''': None,
              '''timesteps''': [2_2, 0],
              '''class_labels''': 0,
              '''generator''': generator,
              '''output_type''': '''np''',
          }

          if get_fixed_latents:
               a__   :      Optional[Any]     =						self.get_fixed_latents(seed=SCREAMING_SNAKE_CASE__      ,						device=SCREAMING_SNAKE_CASE__      ,						dtype=SCREAMING_SNAKE_CASE__      ,						shape=SCREAMING_SNAKE_CASE__	)
               a__   :      List[str]     =						latents

          return inputs







     def 			__lowerCAmelCase	(		self						:	Optional[Any]      ,						A__						:	str=0      ,						A__						:	Union[str, Any]="cpu"      ,						A__						:	List[str]=torch.floataa      ,						A__						:	str=(1, 3, 6_4, 6_4)	)       ->			Optional[int]:
          '''simple docstring'''




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







     def 			__lowerCAmelCase	(		self						:	str	)       ->			Tuple:
          '''simple docstring'''




          a__   :      List[Any]     =						UNetaDModel.from_pretrained('''diffusers/consistency_models'''      ,						subfolder='''diffusers_cd_imagenet64_l2'''	)
          a__   :      int     =						CMStochasticIterativeScheduler(
              num_train_timesteps=4_0      ,						sigma_min=0.002      ,						sigma_max=8_0.0      ,						)
          a__   :      int     =						ConsistencyModelPipeline(unet=SCREAMING_SNAKE_CASE__      ,						scheduler=SCREAMING_SNAKE_CASE__	)
          pipe.to(torch_device=SCREAMING_SNAKE_CASE__	)
          pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__	)

          a__   :      Any     =						self.get_inputs()
          a__   :      List[Any]     =						pipe(**SCREAMING_SNAKE_CASE__	).images
          assert image.shape == (1, 6_4, 6_4, 3)

          a__   :      List[Any]     =						image[0, -3:, -3:, -1]

          a__   :      Dict     =						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 			__lowerCAmelCase	(		self						:	Tuple	)       ->			int:
          '''simple docstring'''




          a__   :      Optional[Any]     =						UNetaDModel.from_pretrained('''diffusers/consistency_models'''      ,						subfolder='''diffusers_cd_imagenet64_l2'''	)
          a__   :      Optional[int]     =						CMStochasticIterativeScheduler(
              num_train_timesteps=4_0      ,						sigma_min=0.002      ,						sigma_max=8_0.0      ,						)
          a__   :      Optional[Any]     =						ConsistencyModelPipeline(unet=SCREAMING_SNAKE_CASE__      ,						scheduler=SCREAMING_SNAKE_CASE__	)
          pipe.to(torch_device=SCREAMING_SNAKE_CASE__	)
          pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__	)

          a__   :      List[Any]     =						self.get_inputs()
          a__   :      Dict     =						1
          a__   :      Optional[Any]     =						None
          a__   :      Optional[Any]     =						pipe(**SCREAMING_SNAKE_CASE__	).images
          assert image.shape == (1, 6_4, 6_4, 3)

          a__   :      Union[str, Any]     =						image[0, -3:, -3:, -1]

          a__   :      Optional[Any]     =						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 			__lowerCAmelCase	(		self						:	Dict	)       ->			Dict:
          '''simple docstring'''




          a__   :      Any     =						UNetaDModel.from_pretrained('''diffusers/consistency_models'''      ,						subfolder='''diffusers_cd_imagenet64_l2'''	)
          a__   :      List[str]     =						CMStochasticIterativeScheduler(
              num_train_timesteps=4_0      ,						sigma_min=0.002      ,						sigma_max=8_0.0      ,						)
          a__   :      List[Any]     =						ConsistencyModelPipeline(unet=SCREAMING_SNAKE_CASE__      ,						scheduler=SCREAMING_SNAKE_CASE__	)
          pipe.to(torch_device=SCREAMING_SNAKE_CASE__      ,						torch_dtype=torch.floataa	)
          pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__	)

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

          a__   :      Union[str, Any]     =						image[0, -3:, -3:, -1]

          a__   :      Any     =						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 			__lowerCAmelCase	(		self						:	int	)       ->			Dict:
          '''simple docstring'''




          a__   :      Optional[int]     =						UNetaDModel.from_pretrained('''diffusers/consistency_models'''      ,						subfolder='''diffusers_cd_imagenet64_l2'''	)
          a__   :      Dict     =						CMStochasticIterativeScheduler(
              num_train_timesteps=4_0      ,						sigma_min=0.002      ,						sigma_max=8_0.0      ,						)
          a__   :      str     =						ConsistencyModelPipeline(unet=SCREAMING_SNAKE_CASE__      ,						scheduler=SCREAMING_SNAKE_CASE__	)
          pipe.to(torch_device=SCREAMING_SNAKE_CASE__      ,						torch_dtype=torch.floataa	)
          pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__	)

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

          a__   :      Optional[int]     =						image[0, -3:, -3:, -1]

          a__   :      Tuple     =						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"""






from collections.abc import Callable

import numpy as np





def 				A_       ( snake_case__       ,					snake_case__       ,					snake_case__       ,					snake_case__       ,					snake_case__		)					->  np.ndarray:
			_UpperCamelCase					:str					     =	int(np.ceil((x_end - xa) / step_size		)		)
			_UpperCamelCase					:Union[str, Any]					     =	np.zeros((n + 1,)		)
			_UpperCamelCase					:List[str]					     =	ya
			_UpperCamelCase					:Any					     =	xa

			for k in range(snake_case__		):
						_UpperCamelCase					:Union[str, Any]					     =	y[k] + step_size * ode_func(snake_case__       ,					y[k]		)
						x += step_size

			return y


if __name__ == "__main__":
					import doctest

					doctest.testmod()

