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






			'''simple docstring'''

			__UpperCamelCase      :						Optional[Any]					      = ConsistencyModelPipeline
			__UpperCamelCase      :						List[str]					      = UNCONDITIONAL_IMAGE_GENERATION_PARAMS
			__UpperCamelCase      :						List[str]					      = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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

			@property
			def       lowerCAmelCase__	(       self :       str      ):
							UpperCamelCase_:							List[str]	=						UNetaDModel.from_pretrained(
							    """diffusers/consistency-models-test"""						,	subfolder="""test_unet"""						,	)
							return unet

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

			def       lowerCAmelCase__	(       self :       Any						,	snake_case_ :       Any=False      ):
							if class_cond:
											UpperCamelCase_:							Optional[int]	=						self.dummy_cond_unet
							else:
											UpperCamelCase_:							Union[str, Any]	=						self.dummy_uncond_unet

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

							UpperCamelCase_:							Tuple	=						{
							    """unet""": unet,
							    """scheduler""": scheduler,
							}

							return components

			def       lowerCAmelCase__	(       self :       Dict						,	snake_case_ :       str						,	snake_case_ :       str=0      ):
							if str(snake_case_      ).startswith("""mps"""      ):
											UpperCamelCase_:							List[Any]	=						torch.manual_seed(snake_case_      )
							else:
											UpperCamelCase_:							int	=						torch.Generator(device=snake_case_      ).manual_seed(snake_case_      )

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

							return inputs

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

							UpperCamelCase_:							Optional[Any]	=						self.get_dummy_inputs(snake_case_      )
							UpperCamelCase_:							Any	=						pipe(**snake_case_      ).images
							assert image.shape == (1, 32, 32, 3)

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

							UpperCamelCase_:							Any	=						self.get_dummy_inputs(snake_case_      )
							UpperCamelCase_:							Tuple	=						0
							UpperCamelCase_:							List[str]	=						pipe(**snake_case_      ).images
							assert image.shape == (1, 32, 32, 3)

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

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

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

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

							UpperCamelCase_:							Any	=						image[0, -3:, -3:, -1]
							UpperCamelCase_:							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  ):






			'''simple docstring'''

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

			def       lowerCAmelCase__	(       self :       Optional[int]						,	snake_case_ :       Any=0						,	snake_case_ :       Tuple=False						,	snake_case_ :       Optional[Any]="cpu"						,	snake_case_ :       str=torch.floataa						,	snake_case_ :       str=(1, 3, 64, 64)      ):
							UpperCamelCase_:							List[Any]	=						torch.manual_seed(snake_case_      )

							UpperCamelCase_:							Dict	=						{
							    """num_inference_steps""": None,
							    """timesteps""": [22, 0],
							    """class_labels""": 0,
							    """generator""": generator,
							    """output_type""": """np""",
							}

							if get_fixed_latents:
											UpperCamelCase_:							List[str]	=						self.get_fixed_latents(seed=snake_case_						,	device=snake_case_						,	dtype=snake_case_						,	shape=snake_case_      )
											UpperCamelCase_:							Any	=						latents

							return inputs

			def       lowerCAmelCase__	(       self :       Any						,	snake_case_ :       Optional[int]=0						,	snake_case_ :       Optional[int]="cpu"						,	snake_case_ :       Any=torch.floataa						,	snake_case_ :       List[Any]=(1, 3, 64, 64)      ):
							if type(snake_case_      ) == str:
											UpperCamelCase_:							Optional[Any]	=						torch.device(snake_case_      )
							UpperCamelCase_:							Optional[Any]	=						torch.Generator(device=snake_case_      ).manual_seed(snake_case_      )
							UpperCamelCase_:							List[str]	=						randn_tensor(snake_case_						,	generator=snake_case_						,	device=snake_case_						,	dtype=snake_case_      )
							return latents

			def       lowerCAmelCase__	(       self :       Any      ):
							UpperCamelCase_:							Any	=						UNetaDModel.from_pretrained("""diffusers/consistency_models"""						,	subfolder="""diffusers_cd_imagenet64_l2"""      )
							UpperCamelCase_:							Tuple	=						CMStochasticIterativeScheduler(
							    num_train_timesteps=40						,	sigma_min=0.002						,	sigma_max=80.0						,	)
							UpperCamelCase_:							List[str]	=						ConsistencyModelPipeline(unet=snake_case_						,	scheduler=snake_case_      )
							pipe.to(torch_device=snake_case_      )
							pipe.set_progress_bar_config(disable=snake_case_      )

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

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

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

							UpperCamelCase_:							List[str]	=						self.get_inputs()
							UpperCamelCase_:							Optional[int]	=						1
							UpperCamelCase_:							str	=						None
							UpperCamelCase_:							List[str]	=						pipe(**snake_case_      ).images
							assert image.shape == (1, 64, 64, 3)

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

							UpperCamelCase_:							int	=						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 :       str      ):
							UpperCamelCase_:							Tuple	=						UNetaDModel.from_pretrained("""diffusers/consistency_models"""						,	subfolder="""diffusers_cd_imagenet64_l2"""      )
							UpperCamelCase_:							List[Any]	=						CMStochasticIterativeScheduler(
							    num_train_timesteps=40						,	sigma_min=0.002						,	sigma_max=80.0						,	)
							UpperCamelCase_:							List[Any]	=						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_      )

							UpperCamelCase_:							Union[str, Any]	=						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_      ):
											UpperCamelCase_:							Optional[int]	=						pipe(**snake_case_      ).images
							assert image.shape == (1, 64, 64, 3)

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

							UpperCamelCase_:							Optional[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 :       List[Any]      ):
							UpperCamelCase_:							Tuple	=						UNetaDModel.from_pretrained("""diffusers/consistency_models"""						,	subfolder="""diffusers_cd_imagenet64_l2"""      )
							UpperCamelCase_:							str	=						CMStochasticIterativeScheduler(
							    num_train_timesteps=40						,	sigma_min=0.002						,	sigma_max=80.0						,	)
							UpperCamelCase_:							List[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_      )

							UpperCamelCase_:							List[str]	=						self.get_inputs(get_fixed_latents=snake_case_						,	device=snake_case_      )
							UpperCamelCase_:							Union[str, Any]	=						1
							UpperCamelCase_:							int	=						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_      ):
											UpperCamelCase_:							Tuple	=						pipe(**snake_case_      ).images
							assert image.shape == (1, 64, 64, 3)

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

							UpperCamelCase_:							Optional[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








import cva
import numpy as np

class        _UpperCamelCase				:






			'''simple docstring'''

			def __init__(       self :       Dict						,	snake_case_ :       float						,	snake_case_ :       int      ):

							if k in (0.04, 0.06):
											UpperCamelCase_:							Union[str, Any]	=						k
											UpperCamelCase_:							Union[str, Any]	=						window_size
							else:
											raise ValueError("""invalid k value"""      )

			def __str__(       self :       int      ):
							return str(self.k      )




			def       lowerCAmelCase__	(       self :       Optional[Any]						,	snake_case_ :       str      ):
							UpperCamelCase_:							int	=						cva.imread(snake_case_						,	0      )
							UpperCamelCase_,   UpperCamelCase_:							List[Any]	=						img.shape
							UpperCamelCase_:							list[list[int]]	=						[]
							UpperCamelCase_:							int	=						img.copy()
							UpperCamelCase_:							Any	=						cva.cvtColor(snake_case_						,	cva.COLOR_GRAY2RGB      )
							UpperCamelCase_,   UpperCamelCase_:							List[Any]	=						np.gradient(snake_case_      )
							UpperCamelCase_:							Optional[Any]	=						dx**2
							UpperCamelCase_:							Dict	=						dy**2
							UpperCamelCase_:							Optional[Any]	=						dx * dy
							UpperCamelCase_:							str	=						0.04
							UpperCamelCase_:							int	=						self.window_size // 2
							for y in range(snake_case_						,	h - offset      ):
											for x in range(snake_case_						,	w - offset      ):
															UpperCamelCase_:							List[Any]	=						ixx[
															    y - offset : y + offset + 1, x - offset : x + offset + 1
															].sum()
															UpperCamelCase_:							int	=						iyy[
															    y - offset : y + offset + 1, x - offset : x + offset + 1
															].sum()
															UpperCamelCase_:							List[str]	=						ixy[
															    y - offset : y + offset + 1, x - offset : x + offset + 1
															].sum()

															UpperCamelCase_:							List[str]	=						(wxx * wyy) - (wxy**2)
															UpperCamelCase_:							Optional[int]	=						wxx + wyy
															UpperCamelCase_:							Dict	=						det - k * (trace**2)
															# Can change the value
															if r > 0.5:
																			corner_list.append([x, y, r]      )
																			color_img.itemset((y, x, 0)						,	0      )
																			color_img.itemset((y, x, 1)						,	0      )
																			color_img.itemset((y, x, 2)						,	255      )
							return color_img, corner_list


if __name__ == "__main__":
							lowerCamelCase_						: Optional[Any]         =  HarrisCorner(0.04, 3)
							lowerCamelCase_   ,       lowerCamelCase_						: Any         =  edge_detect.detect("""path_to_image""")
							cva.imwrite("""detect.png""", color_img)






