
"""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 				A__ (			UpperCAmelCase_ ,  unittest.TestCase):
		A_		: List[str]            =			ConsistencyModelPipeline
		A_		: int            =			UNCONDITIONAL_IMAGE_GENERATION_PARAMS
		A_		: Tuple            =			UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

		# Override required_optional_params to remove num_images_per_prompt
		A_		: Dict            =			frozenset(
		    [
		        'num_inference_steps',
		        'generator',
		        'latents',
		        'output_type',
		        'return_dict',
		        'callback',
		        'callback_steps',
		    ])

		@property
		def 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		Any		    =   UNetaDModel.from_pretrained(
			    'diffusers/consistency-models-test' ,	subfolder='test_unet' ,	)
			return unet

		@property
		def 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		int		    =   UNetaDModel.from_pretrained(
			    'diffusers/consistency-models-test' ,	subfolder='test_unet_class_cond' ,	)
			return unet

		def 							__lowerCamelCase							(   self ,	_SCREAMING_SNAKE_CASE=False				):
			if class_cond:
				__lowerCAmelCase    :		Union[str, Any]		    =   self.dummy_cond_unet
			else:
				__lowerCAmelCase    :		str		    =   self.dummy_uncond_unet

			# Default to CM multistep sampler
			__lowerCAmelCase    :		Union[str, Any]		    =   CMStochasticIterativeScheduler(
			    num_train_timesteps=40 ,	sigma_min=0.002 ,	sigma_max=80.0 ,	)

			__lowerCAmelCase    :		Any		    =   {
			    'unet': unet,
			    'scheduler': scheduler,
			}

			return components

		def 							__lowerCamelCase							(   self ,	_SCREAMING_SNAKE_CASE ,	_SCREAMING_SNAKE_CASE=0				):
			if str(_lowercase				).startswith('mps'				):
				__lowerCAmelCase    :		Union[str, Any]		    =   torch.manual_seed(_lowercase				)
			else:
				__lowerCAmelCase    :		int		    =   torch.Generator(device=_lowercase				).manual_seed(_lowercase				)

			__lowerCAmelCase    :		str		    =   {
			    'batch_size': 1,
			    'num_inference_steps': None,
			    'timesteps': [22, 0],
			    'generator': generator,
			    'output_type': 'np',
			}

			return inputs

		def 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		List[Any]		    =   'cpu'  # ensure determinism for the device-dependent torch.Generator
			__lowerCAmelCase    :		Dict		    =   self.get_dummy_components()
			__lowerCAmelCase    :		int		    =   ConsistencyModelPipeline(**_lowercase				)
			__lowerCAmelCase    :		str		    =   pipe.to(_lowercase				)
			pipe.set_progress_bar_config(disable=_lowercase				)

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

			__lowerCAmelCase    :		Tuple		    =   image[0, -3:, -3:, -1]
			__lowerCAmelCase    :		List[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 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		Tuple		    =   'cpu'  # ensure determinism for the device-dependent torch.Generator
			__lowerCAmelCase    :		Any		    =   self.get_dummy_components(class_cond=_lowercase				)
			__lowerCAmelCase    :		List[str]		    =   ConsistencyModelPipeline(**_lowercase				)
			__lowerCAmelCase    :		Tuple		    =   pipe.to(_lowercase				)
			pipe.set_progress_bar_config(disable=_lowercase				)

			__lowerCAmelCase    :		List[str]		    =   self.get_dummy_inputs(_lowercase				)
			__lowerCAmelCase    :		Dict		    =   0
			__lowerCAmelCase    :		Tuple		    =   pipe(**_lowercase				).images
			assert image.shape == (1, 32, 32, 3)

			__lowerCAmelCase    :		Tuple		    =   image[0, -3:, -3:, -1]
			__lowerCAmelCase    :		str		    =   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 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		int		    =   'cpu'  # ensure determinism for the device-dependent torch.Generator
			__lowerCAmelCase    :		Dict		    =   self.get_dummy_components()
			__lowerCAmelCase    :		List[str]		    =   ConsistencyModelPipeline(**_lowercase				)
			__lowerCAmelCase    :		Union[str, Any]		    =   pipe.to(_lowercase				)
			pipe.set_progress_bar_config(disable=_lowercase				)

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

			__lowerCAmelCase    :		int		    =   image[0, -3:, -3:, -1]
			__lowerCAmelCase    :		Optional[Any]		    =   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 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		str		    =   'cpu'  # ensure determinism for the device-dependent torch.Generator
			__lowerCAmelCase    :		Tuple		    =   self.get_dummy_components(class_cond=_lowercase				)
			__lowerCAmelCase    :		Any		    =   ConsistencyModelPipeline(**_lowercase				)
			__lowerCAmelCase    :		Any		    =   pipe.to(_lowercase				)
			pipe.set_progress_bar_config(disable=_lowercase				)

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

			__lowerCAmelCase    :		Tuple		    =   image[0, -3:, -3:, -1]
			__lowerCAmelCase    :		Tuple		    =   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 				A__ (			unittest.TestCase):

		def 							__lowerCamelCase							(   self				):
			super().tearDown()
			gc.collect()
			torch.cuda.empty_cache()

		def 							__lowerCamelCase							(   self ,	_SCREAMING_SNAKE_CASE=0 ,	_SCREAMING_SNAKE_CASE=False ,	_SCREAMING_SNAKE_CASE="cpu" ,	_SCREAMING_SNAKE_CASE=torch.floataa ,	_SCREAMING_SNAKE_CASE=(1, 3, 64, 64)				):
			__lowerCAmelCase    :		List[str]		    =   torch.manual_seed(_lowercase				)

			__lowerCAmelCase    :		int		    =   {
			    'num_inference_steps': None,
			    'timesteps': [22, 0],
			    'class_labels': 0,
			    'generator': generator,
			    'output_type': 'np',
			}

			if get_fixed_latents:
				__lowerCAmelCase    :		Any		    =   self.get_fixed_latents(seed=_lowercase ,	device=_lowercase ,	dtype=_lowercase ,	shape=_lowercase				)
				__lowerCAmelCase    :		Union[str, Any]		    =   latents

			return inputs

		def 							__lowerCamelCase							(   self ,	_SCREAMING_SNAKE_CASE=0 ,	_SCREAMING_SNAKE_CASE="cpu" ,	_SCREAMING_SNAKE_CASE=torch.floataa ,	_SCREAMING_SNAKE_CASE=(1, 3, 64, 64)				):
			if type(_lowercase				) == str:
				__lowerCAmelCase    :		Dict		    =   torch.device(_lowercase				)
			__lowerCAmelCase    :		List[str]		    =   torch.Generator(device=_lowercase				).manual_seed(_lowercase				)
			__lowerCAmelCase    :		List[Any]		    =   randn_tensor(_lowercase ,	generator=_lowercase ,	device=_lowercase ,	dtype=_lowercase				)
			return latents

		def 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		List[str]		    =   UNetaDModel.from_pretrained('diffusers/consistency_models' ,	subfolder='diffusers_cd_imagenet64_l2'				)
			__lowerCAmelCase    :		Any		    =   CMStochasticIterativeScheduler(
			    num_train_timesteps=40 ,	sigma_min=0.002 ,	sigma_max=80.0 ,	)
			__lowerCAmelCase    :		List[str]		    =   ConsistencyModelPipeline(unet=_lowercase ,	scheduler=_lowercase				)
			pipe.to(torch_device=_lowercase				)
			pipe.set_progress_bar_config(disable=_lowercase				)

			__lowerCAmelCase    :		List[str]		    =   self.get_inputs()
			__lowerCAmelCase    :		int		    =   pipe(**_lowercase				).images
			assert image.shape == (1, 64, 64, 3)

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

			__lowerCAmelCase    :		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 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		List[str]		    =   UNetaDModel.from_pretrained('diffusers/consistency_models' ,	subfolder='diffusers_cd_imagenet64_l2'				)
			__lowerCAmelCase    :		List[str]		    =   CMStochasticIterativeScheduler(
			    num_train_timesteps=40 ,	sigma_min=0.002 ,	sigma_max=80.0 ,	)
			__lowerCAmelCase    :		Tuple		    =   ConsistencyModelPipeline(unet=_lowercase ,	scheduler=_lowercase				)
			pipe.to(torch_device=_lowercase				)
			pipe.set_progress_bar_config(disable=_lowercase				)

			__lowerCAmelCase    :		Tuple		    =   self.get_inputs()
			__lowerCAmelCase    :		Any		    =   1
			__lowerCAmelCase    :		int		    =   None
			__lowerCAmelCase    :		Optional[int]		    =   pipe(**_lowercase				).images
			assert image.shape == (1, 64, 64, 3)

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

			__lowerCAmelCase    :		Tuple		    =   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 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		Optional[int]		    =   UNetaDModel.from_pretrained('diffusers/consistency_models' ,	subfolder='diffusers_cd_imagenet64_l2'				)
			__lowerCAmelCase    :		Optional[Any]		    =   CMStochasticIterativeScheduler(
			    num_train_timesteps=40 ,	sigma_min=0.002 ,	sigma_max=80.0 ,	)
			__lowerCAmelCase    :		Tuple		    =   ConsistencyModelPipeline(unet=_lowercase ,	scheduler=_lowercase				)
			pipe.to(torch_device=_lowercase ,	torch_dtype=torch.floataa				)
			pipe.set_progress_bar_config(disable=_lowercase				)

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

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

			__lowerCAmelCase    :		int		    =   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 							__lowerCamelCase							(   self				):
			__lowerCAmelCase    :		List[Any]		    =   UNetaDModel.from_pretrained('diffusers/consistency_models' ,	subfolder='diffusers_cd_imagenet64_l2'				)
			__lowerCAmelCase    :		Any		    =   CMStochasticIterativeScheduler(
			    num_train_timesteps=40 ,	sigma_min=0.002 ,	sigma_max=80.0 ,	)
			__lowerCAmelCase    :		List[str]		    =   ConsistencyModelPipeline(unet=_lowercase ,	scheduler=_lowercase				)
			pipe.to(torch_device=_lowercase ,	torch_dtype=torch.floataa				)
			pipe.set_progress_bar_config(disable=_lowercase				)

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

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

			__lowerCAmelCase    :		List[str]		    =   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

"""simple docstring"""







from .testing import (
    are_the_same_tensors,
    execute_subprocess_async,
    require_bnb,
    require_cpu,
    require_cuda,
    require_huggingface_suite,
    require_mps,
    require_multi_gpu,
    require_multi_xpu,
    require_safetensors,
    require_single_gpu,
    require_single_xpu,
    require_torch_min_version,
    require_tpu,
    require_xpu,
    skip,
    slow,
)
from .training import RegressionDataset, RegressionModel, RegressionModelaXPU


from .scripts import test_script, test_sync, test_ops  # isort: skip