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




	'''simple docstring'''




	_A								=	ConsistencyModelPipeline
	_A								=	UNCONDITIONAL_IMAGE_GENERATION_PARAMS
	_A								=	UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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







	@property
	def   _lowerCamelCase							(   self						:Dict		)						->					Tuple:
								__UpperCamelCase      :  str       =   UNetaDModel.from_pretrained(
								    "diffusers/consistency-models-test"	,  subfolder="test_unet"	,  )
								return unet







	@property
	def   _lowerCamelCase							(   self						:List[Any]		)						->					Optional[int]:
								__UpperCamelCase      :  List[Any]       =   UNetaDModel.from_pretrained(
								    "diffusers/consistency-models-test"	,  subfolder="test_unet_class_cond"	,  )
								return unet







	def   _lowerCamelCase							(   self						:Union[str, Any]	,  a						:List[Any]=False		)						->					Optional[int]:
								if class_cond:
															__UpperCamelCase      :  str       =   self.dummy_cond_unet
								else:
															__UpperCamelCase      :  Dict       =   self.dummy_uncond_unet

								# Default to CM multistep sampler
								__UpperCamelCase      :  Optional[Any]       =   CMStochasticIterativeScheduler(
								    num_train_timesteps=4_0	,  sigma_min=0.002	,  sigma_max=80.0	,  )

								__UpperCamelCase      :  Optional[int]       =   {
								    "unet": unet,
								    "scheduler": scheduler,
								}

								return components







	def   _lowerCamelCase							(   self						:Optional[Any]	,  a						:Tuple	,  a						:Optional[int]=0		)						->					int:
								if str(a		).startswith("mps"		):
															__UpperCamelCase      :  Dict       =   torch.manual_seed(a		)
								else:
															__UpperCamelCase      :  Dict       =   torch.Generator(device=a		).manual_seed(a		)

								__UpperCamelCase      :  str       =   {
								    "batch_size": 1,
								    "num_inference_steps": None,
								    "timesteps": [2_2, 0],
								    "generator": generator,
								    "output_type": "np",
								}

								return inputs







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

								__UpperCamelCase      :  Optional[int]       =   self.get_dummy_inputs(a		)
								__UpperCamelCase      :  str       =   pipe(**a		).images
								assert image.shape == (1, 3_2, 3_2, 3)

								__UpperCamelCase      :  Optional[Any]       =   image[0, -3:, -3:, -1]
								__UpperCamelCase      :  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   _lowerCamelCase							(   self						:Optional[Any]		)						->					Optional[Any]:
								__UpperCamelCase      :  Optional[Any]       =   "cpu"  # ensure determinism for the device-dependent torch.Generator
								__UpperCamelCase      :  List[str]       =   self.get_dummy_components(class_cond=a		)
								__UpperCamelCase      :  Union[str, Any]       =   ConsistencyModelPipeline(**a		)
								__UpperCamelCase      :  List[str]       =   pipe.to(a		)
								pipe.set_progress_bar_config(disable=a		)

								__UpperCamelCase      :  List[Any]       =   self.get_dummy_inputs(a		)
								__UpperCamelCase      :  str       =   0
								__UpperCamelCase      :  str       =   pipe(**a		).images
								assert image.shape == (1, 3_2, 3_2, 3)

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

								__UpperCamelCase      :  Any       =   self.get_dummy_inputs(a		)
								__UpperCamelCase      :  Optional[int]       =   1
								__UpperCamelCase      :  Tuple       =   None
								__UpperCamelCase      :  Union[str, Any]       =   pipe(**a		).images
								assert image.shape == (1, 3_2, 3_2, 3)

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

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

								__UpperCamelCase      :  int       =   image[0, -3:, -3:, -1]
								__UpperCamelCase      :  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




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




	'''simple docstring'''




	def   _lowerCamelCase							(   self						:Any		)						->					Dict:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







	def   _lowerCamelCase							(   self						:Any	,  a						:List[str]=0	,  a						:Any=False	,  a						:Any="cpu"	,  a						:List[str]=torch.floataa	,  a						:List[str]=(1, 3, 6_4, 6_4)		)						->					Optional[int]:
								__UpperCamelCase      :  Tuple       =   torch.manual_seed(a		)

								__UpperCamelCase      :  Union[str, Any]       =   {
								    "num_inference_steps": None,
								    "timesteps": [2_2, 0],
								    "class_labels": 0,
								    "generator": generator,
								    "output_type": "np",
								}

								if get_fixed_latents:
															__UpperCamelCase      :  Optional[int]       =   self.get_fixed_latents(seed=a	,  device=a	,  dtype=a	,  shape=a		)
															__UpperCamelCase      :  str       =   latents

								return inputs







	def   _lowerCamelCase							(   self						:Optional[int]	,  a						:Tuple=0	,  a						:str="cpu"	,  a						:Optional[int]=torch.floataa	,  a						:Optional[Any]=(1, 3, 6_4, 6_4)		)						->					Optional[Any]:
								if type(a		) == str:
															__UpperCamelCase      :  int       =   torch.device(a		)
								__UpperCamelCase      :  Dict       =   torch.Generator(device=a		).manual_seed(a		)
								__UpperCamelCase      :  Optional[Any]       =   randn_tensor(a	,  generator=a	,  device=a	,  dtype=a		)
								return latents







	def   _lowerCamelCase							(   self						:List[str]		)						->					Union[str, Any]:
								__UpperCamelCase      :  Optional[Any]       =   UNetaDModel.from_pretrained("diffusers/consistency_models"	,  subfolder="diffusers_cd_imagenet64_l2"		)
								__UpperCamelCase      :  int       =   CMStochasticIterativeScheduler(
								    num_train_timesteps=4_0	,  sigma_min=0.002	,  sigma_max=80.0	,  )
								__UpperCamelCase      :  List[str]       =   ConsistencyModelPipeline(unet=a	,  scheduler=a		)
								pipe.to(torch_device=a		)
								pipe.set_progress_bar_config(disable=a		)

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

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

								__UpperCamelCase      :  Any       =   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						:Any		)						->					Tuple:
								__UpperCamelCase      :  int       =   UNetaDModel.from_pretrained("diffusers/consistency_models"	,  subfolder="diffusers_cd_imagenet64_l2"		)
								__UpperCamelCase      :  List[str]       =   CMStochasticIterativeScheduler(
								    num_train_timesteps=4_0	,  sigma_min=0.002	,  sigma_max=80.0	,  )
								__UpperCamelCase      :  str       =   ConsistencyModelPipeline(unet=a	,  scheduler=a		)
								pipe.to(torch_device=a		)
								pipe.set_progress_bar_config(disable=a		)

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

								__UpperCamelCase      :  List[str]       =   image[0, -3:, -3:, -1]

								__UpperCamelCase      :  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   _lowerCamelCase							(   self						:Any		)						->					Any:
								__UpperCamelCase      :  List[str]       =   UNetaDModel.from_pretrained("diffusers/consistency_models"	,  subfolder="diffusers_cd_imagenet64_l2"		)
								__UpperCamelCase      :  List[Any]       =   CMStochasticIterativeScheduler(
								    num_train_timesteps=4_0	,  sigma_min=0.002	,  sigma_max=80.0	,  )
								__UpperCamelCase      :  Tuple       =   ConsistencyModelPipeline(unet=a	,  scheduler=a		)
								pipe.to(torch_device=a	,  torch_dtype=torch.floataa		)
								pipe.set_progress_bar_config(disable=a		)

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

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

								__UpperCamelCase      :  str       =   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						:List[str]		)						->					Dict:
								__UpperCamelCase      :  Optional[Any]       =   UNetaDModel.from_pretrained("diffusers/consistency_models"	,  subfolder="diffusers_cd_imagenet64_l2"		)
								__UpperCamelCase      :  Optional[int]       =   CMStochasticIterativeScheduler(
								    num_train_timesteps=4_0	,  sigma_min=0.002	,  sigma_max=80.0	,  )
								__UpperCamelCase      :  List[str]       =   ConsistencyModelPipeline(unet=a	,  scheduler=a		)
								pipe.to(torch_device=a	,  torch_dtype=torch.floataa		)
								pipe.set_progress_bar_config(disable=a		)

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

								__UpperCamelCase      :  str       =   image[0, -3:, -3:, -1]

								__UpperCamelCase      :  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

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


lowercase    :   int						=			{
    'configuration_conditional_detr': [
        'CONDITIONAL_DETR_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'ConditionalDetrConfig',
        'ConditionalDetrOnnxConfig',
    ]
}

try:
		if not is_vision_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		lowercase    :   List[str]						=			['ConditionalDetrFeatureExtractor']
		lowercase    :   int						=			['ConditionalDetrImageProcessor']

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		lowercase    :   Optional[int]						=			[
		    'CONDITIONAL_DETR_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'ConditionalDetrForObjectDetection',
		    'ConditionalDetrForSegmentation',
		    'ConditionalDetrModel',
		    'ConditionalDetrPreTrainedModel',
		]


if TYPE_CHECKING:
		from .configuration_conditional_detr import (
		    CONDITIONAL_DETR_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    ConditionalDetrConfig,
		    ConditionalDetrOnnxConfig,
		)

		try:
				if not is_vision_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .feature_extraction_conditional_detr import ConditionalDetrFeatureExtractor
				from .image_processing_conditional_detr import ConditionalDetrImageProcessor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_conditional_detr import (
				    CONDITIONAL_DETR_PRETRAINED_MODEL_ARCHIVE_LIST,
				    ConditionalDetrForObjectDetection,
				    ConditionalDetrForSegmentation,
				    ConditionalDetrModel,
				    ConditionalDetrPreTrainedModel,
				)

else:
		import sys

		lowercase    :   Optional[int]						=			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)