






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 							snake_case_       (    __UpperCamelCase		,							unittest.TestCase      ):

						"""simple docstring"""



						snake_case__        =      ConsistencyModelPipeline
						snake_case__        =      UNCONDITIONAL_IMAGE_GENERATION_PARAMS
						snake_case__        =      UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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



						@property
						def        UpperCAmelCase__					(self:					Any   )    ->     Any:




											'''simple docstring'''



											__a					:						Optional[Any]      =					UNetaDModel.from_pretrained(
											    "diffusers/consistency-models-test" ,					subfolder="test_unet" ,					)
											return unet



						@property
						def        UpperCAmelCase__					(self:					int   )    ->     Union[str, Any]:




											'''simple docstring'''



											__a					:						int      =					UNetaDModel.from_pretrained(
											    "diffusers/consistency-models-test" ,					subfolder="test_unet_class_cond" ,					)
											return unet



						def        UpperCAmelCase__					(self:					Tuple ,					__UpperCAmelCase:					Union[str, Any]=False   )    ->     int:




											'''simple docstring'''



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

											# Default to CM multistep sampler
											__a					:						Optional[int]      =					CMStochasticIterativeScheduler(
											    num_train_timesteps=40 ,					sigma_min=0.0_02 ,					sigma_max=80.0 ,					)

											__a					:						int      =					{
											    "unet": unet,
											    "scheduler": scheduler,
											}

											return components



						def        UpperCAmelCase__					(self:					int ,					__UpperCAmelCase:					Optional[int] ,					__UpperCAmelCase:					Tuple=0   )    ->     Optional[Any]:




											'''simple docstring'''



											if str(__UpperCAmelCase   ).startswith("mps"   ):
																__a					:						Tuple      =					torch.manual_seed(__UpperCAmelCase   )
											else:
																__a					:						List[str]      =					torch.Generator(device=__UpperCAmelCase   ).manual_seed(__UpperCAmelCase   )

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

											return inputs



						def        UpperCAmelCase__					(self:					Optional[int]   )    ->     Optional[int]:




											'''simple docstring'''



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

											__a					:						Any      =					self.get_dummy_inputs(__UpperCAmelCase   )
											__a					:						Union[str, Any]      =					pipe(**__UpperCAmelCase   ).images
											assert image.shape == (1, 32, 32, 3)

											__a					:						Any      =					image[0, -3:, -3:, -1]
											__a					:						Optional[Any]      =					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:					List[Any]   )    ->     Optional[int]:




											'''simple docstring'''



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

											__a					:						Optional[int]      =					self.get_dummy_inputs(__UpperCAmelCase   )
											__a					:						Tuple      =					0
											__a					:						List[Any]      =					pipe(**__UpperCAmelCase   ).images
											assert image.shape == (1, 32, 32, 3)

											__a					:						str      =					image[0, -3:, -3:, -1]
											__a					:						Union[str, Any]      =					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:					Optional[int]   )    ->     Optional[Any]:




											'''simple docstring'''



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

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

											__a					:						List[str]      =					image[0, -3:, -3:, -1]
											__a					:						int      =					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:					Union[str, Any]   )    ->     Any:




											'''simple docstring'''



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

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

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

						"""simple docstring"""



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




											'''simple docstring'''



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



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




											'''simple docstring'''



											__a					:						Tuple      =					torch.manual_seed(__UpperCAmelCase   )

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

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

											return inputs



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




											'''simple docstring'''



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



						def        UpperCAmelCase__					(self:					Optional[Any]   )    ->     str:




											'''simple docstring'''



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

											__a					:						Tuple      =					self.get_inputs()
											__a					:						Union[str, Any]      =					pipe(**__UpperCAmelCase   ).images
											assert image.shape == (1, 64, 64, 3)

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

											__a					:						Optional[int]      =					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:					Union[str, Any]   )    ->     List[Any]:




											'''simple docstring'''



											__a					:						str      =					UNetaDModel.from_pretrained("diffusers/consistency_models" ,					subfolder="diffusers_cd_imagenet64_l2"   )
											__a					:						Union[str, Any]      =					CMStochasticIterativeScheduler(
											    num_train_timesteps=40 ,					sigma_min=0.0_02 ,					sigma_max=80.0 ,					)
											__a					:						int      =					ConsistencyModelPipeline(unet=__UpperCAmelCase ,					scheduler=__UpperCAmelCase   )
											pipe.to(torch_device=__UpperCAmelCase   )
											pipe.set_progress_bar_config(disable=__UpperCAmelCase   )

											__a					:						List[Any]      =					self.get_inputs()
											__a					:						Union[str, Any]      =					1
											__a					:						Dict      =					None
											__a					:						int      =					pipe(**__UpperCAmelCase   ).images
											assert image.shape == (1, 64, 64, 3)

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

											__a					:						int      =					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:					str   )    ->     Optional[Any]:




											'''simple docstring'''



											__a					:						List[Any]      =					UNetaDModel.from_pretrained("diffusers/consistency_models" ,					subfolder="diffusers_cd_imagenet64_l2"   )
											__a					:						Tuple      =					CMStochasticIterativeScheduler(
											    num_train_timesteps=40 ,					sigma_min=0.0_02 ,					sigma_max=80.0 ,					)
											__a					:						List[str]      =					ConsistencyModelPipeline(unet=__UpperCAmelCase ,					scheduler=__UpperCAmelCase   )
											pipe.to(torch_device=__UpperCAmelCase ,					torch_dtype=torch.floataa   )
											pipe.set_progress_bar_config(disable=__UpperCAmelCase   )

											__a					:						Optional[Any]      =					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   ):
																__a					:						List[str]      =					pipe(**__UpperCAmelCase   ).images
											assert image.shape == (1, 64, 64, 3)

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

											__a					:						Union[str, Any]      =					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:					List[Any]   )    ->     str:




											'''simple docstring'''



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

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

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

											__a					:						Optional[int]      =					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













from abc import ABC, abstractmethod
from argparse import ArgumentParser




class 							snake_case_       (    __UpperCamelCase      ):

						"""simple docstring"""



						@staticmethod
						@abstractmethod
						def        UpperCAmelCase__					(__UpperCAmelCase:					ArgumentParser   )    ->     Tuple:




											'''simple docstring'''



											raise NotImplementedError()



						@abstractmethod
						def        UpperCAmelCase__					(self:					List[str]   )    ->     List[str]:




											'''simple docstring'''



											raise NotImplementedError()





