






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     _SCREAMING_SNAKE_CASE       (			_a		,   unittest.TestCase      ):
	snake_case__							:     List[Any]     	=  ConsistencyModelPipeline
	snake_case__							:     List[str]     	=  UNCONDITIONAL_IMAGE_GENERATION_PARAMS
	snake_case__							:     int     	=  UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS

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





	@property
	def 						_A (		self				:					Any	):
		UpperCamelCase     :Dict          =	UNetaDModel.from_pretrained(
		    """diffusers/consistency-models-test""" ,  subfolder="""test_unet""" ,  )
		return unet





	@property
	def 						_A (		self				:					Optional[Any]	):
		UpperCamelCase     :str          =	UNetaDModel.from_pretrained(
		    """diffusers/consistency-models-test""" ,  subfolder="""test_unet_class_cond""" ,  )
		return unet





	def 						_A (		self				:					List[str] ,  __lowerCamelCase				:					Any=False	):
		if class_cond:
			UpperCamelCase     :Optional[int]          =	self.dummy_cond_unet
		else:
			UpperCamelCase     :List[str]          =	self.dummy_uncond_unet

		# Default to CM multistep sampler
		UpperCamelCase     :str          =	CMStochasticIterativeScheduler(
		    num_train_timesteps=40 ,  sigma_min=0.002 ,  sigma_max=80.0 ,  )

		UpperCamelCase     :Dict          =	{
		    """unet""": unet,
		    """scheduler""": scheduler,
		}

		return components





	def 						_A (		self				:					Dict ,  __lowerCamelCase				:					Tuple ,  __lowerCamelCase				:					str=0	):
		if str(__lowerCamelCase	).startswith("""mps"""	):
			UpperCamelCase     :Optional[Any]          =	torch.manual_seed(__lowerCamelCase	)
		else:
			UpperCamelCase     :int          =	torch.Generator(device=__lowerCamelCase	).manual_seed(__lowerCamelCase	)

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

		return inputs





	def 						_A (		self				:					Union[str, Any]	):
		UpperCamelCase     :Union[str, Any]          =	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
		UpperCamelCase     :Optional[int]          =	self.get_dummy_components()
		UpperCamelCase     :Tuple          =	ConsistencyModelPipeline(**__lowerCamelCase	)
		UpperCamelCase     :int          =	pipe.to(__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

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

		UpperCamelCase     :str          =	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 						_A (		self				:					int	):
		UpperCamelCase     :List[str]          =	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
		UpperCamelCase     :List[Any]          =	self.get_dummy_components(class_cond=__lowerCamelCase	)
		UpperCamelCase     :List[str]          =	ConsistencyModelPipeline(**__lowerCamelCase	)
		UpperCamelCase     :List[str]          =	pipe.to(__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

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

		UpperCamelCase     :Dict          =	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 						_A (		self				:					int	):
		UpperCamelCase     :List[str]          =	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
		UpperCamelCase     :Dict          =	self.get_dummy_components()
		UpperCamelCase     :int          =	ConsistencyModelPipeline(**__lowerCamelCase	)
		UpperCamelCase     :List[str]          =	pipe.to(__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

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

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

		UpperCamelCase     :Tuple          =	self.get_dummy_inputs(__lowerCamelCase	)
		UpperCamelCase     :Dict          =	1
		UpperCamelCase     :str          =	None
		UpperCamelCase     :int          =	0
		UpperCamelCase     :Tuple          =	pipe(**__lowerCamelCase	).images
		assert image.shape == (1, 32, 32, 3)

		UpperCamelCase     :List[str]          =	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     _SCREAMING_SNAKE_CASE       (			unittest.TestCase      ):





	def 						_A (		self				:					List[str]	):
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()





	def 						_A (		self				:					Tuple ,  __lowerCamelCase				:					List[str]=0 ,  __lowerCamelCase				:					List[Any]=False ,  __lowerCamelCase				:					Tuple="cpu" ,  __lowerCamelCase				:					Dict=torch.floataa ,  __lowerCamelCase				:					Optional[int]=(1, 3, 64, 64)	):
		UpperCamelCase     :str          =	torch.manual_seed(__lowerCamelCase	)

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

		if get_fixed_latents:
			UpperCamelCase     :Optional[Any]          =	self.get_fixed_latents(seed=__lowerCamelCase ,  device=__lowerCamelCase ,  dtype=__lowerCamelCase ,  shape=__lowerCamelCase	)
			UpperCamelCase     :Tuple          =	latents

		return inputs





	def 						_A (		self				:					Tuple ,  __lowerCamelCase				:					List[Any]=0 ,  __lowerCamelCase				:					Union[str, Any]="cpu" ,  __lowerCamelCase				:					str=torch.floataa ,  __lowerCamelCase				:					List[str]=(1, 3, 64, 64)	):
		if type(__lowerCamelCase	) == str:
			UpperCamelCase     :Tuple          =	torch.device(__lowerCamelCase	)
		UpperCamelCase     :List[Any]          =	torch.Generator(device=__lowerCamelCase	).manual_seed(__lowerCamelCase	)
		UpperCamelCase     :Optional[Any]          =	randn_tensor(__lowerCamelCase ,  generator=__lowerCamelCase ,  device=__lowerCamelCase ,  dtype=__lowerCamelCase	)
		return latents





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

		UpperCamelCase     :List[str]          =	self.get_inputs()
		UpperCamelCase     :Dict          =	pipe(**__lowerCamelCase	).images
		assert image.shape == (1, 64, 64, 3)

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

		UpperCamelCase     :List[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 						_A (		self				:					Union[str, Any]	):
		UpperCamelCase     :Tuple          =	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     :Tuple          =	ConsistencyModelPipeline(unet=__lowerCamelCase ,  scheduler=__lowerCamelCase	)
		pipe.to(torch_device=__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

		UpperCamelCase     :List[str]          =	self.get_inputs()
		UpperCamelCase     :Optional[Any]          =	1
		UpperCamelCase     :List[Any]          =	None
		UpperCamelCase     :Any          =	pipe(**__lowerCamelCase	).images
		assert image.shape == (1, 64, 64, 3)

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

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

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

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

		UpperCamelCase     :Dict          =	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 						_A (		self				:					List[Any]	):
		UpperCamelCase     :Any          =	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     :Union[str, Any]          =	ConsistencyModelPipeline(unet=__lowerCamelCase ,  scheduler=__lowerCamelCase	)
		pipe.to(torch_device=__lowerCamelCase ,  torch_dtype=torch.floataa	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

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

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

		UpperCamelCase     :Dict          =	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 collections.abc import Generator
from math import sin

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  bytes					)		->			bytes:






	"""simple docstring"""



	if len(__magic_name__					) != 32:
		raise ValueError("""Input must be of length 32"""					)

	UpperCamelCase     :int          =	B""""""
	for i in [3, 2, 1, 0]:
		little_endian += string_aa[8 * i : 8 * i + 8]
	return little_endian

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  int					)		->			bytes:






	"""simple docstring"""



	if i < 0:
		raise ValueError("""Input must be non-negative"""					)

	UpperCamelCase     :Any          =	format(__magic_name__	,			"""08x"""					)[-8:]
	UpperCamelCase     :Union[str, Any]          =	B""""""
	for i in [3, 2, 1, 0]:
		little_endian_hex += hex_rep[2 * i : 2 * i + 2].encode("""utf-8"""					)
	return little_endian_hex

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  bytes					)		->			bytes:






	"""simple docstring"""



	UpperCamelCase     :str          =	B""""""
	for char in message:
		bit_string += format(__magic_name__	,			"""08b"""					).encode("""utf-8"""					)
	UpperCamelCase     :Any          =	format(len(__magic_name__					)	,			"""064b"""					).encode("""utf-8"""					)

	# Pad bit_string to a multiple of 512 chars
	bit_string += b"1"
	while len(__magic_name__					) % 512 != 448:
		bit_string += b"0"
	bit_string += to_little_endian(start_len[32:]					) + to_little_endian(start_len[:32]					)

	return bit_string

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  bytes					)		->			Generator[list[int], None, None]:






	"""simple docstring"""



	if len(__magic_name__					) % 512 != 0:
		raise ValueError("""Input must have length that's a multiple of 512"""					)

	for pos in range(0	,			len(__magic_name__					)	,			512					):
		UpperCamelCase     :Tuple          =	bit_string[pos : pos + 512]
		UpperCamelCase     :Optional[int]          =	[]
		for i in range(0	,			512	,			32					):
			block_words.append(int(to_little_endian(block[i : i + 32]					)	,			2					)					)
		yield block_words

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  int					)		->			int:






	"""simple docstring"""



	if i < 0:
		raise ValueError("""Input must be non-negative"""					)

	UpperCamelCase     :List[str]          =	format(__magic_name__	,			"""032b"""					)
	UpperCamelCase     :Any          =	""""""
	for c in i_str:
		new_str += "1" if c == "0" else "0"
	return int(__magic_name__	,			2					)

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  int	,			__magic_name__					:  int					)		->			int:






	"""simple docstring"""



	return (a + b) % 2**32

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  int	,			__magic_name__					:  int					)		->			int:






	"""simple docstring"""



	if i < 0:
		raise ValueError("""Input must be non-negative"""					)
	if shift < 0:
		raise ValueError("""Shift must be non-negative"""					)
	return ((i << shift) ^ (i >> (32 - shift))) % 2**32

def    SCREAMING_SNAKE_CASE_							(							__magic_name__					:  bytes					)		->			bytes:






	"""simple docstring"""



	UpperCamelCase     :Tuple          =	preprocess(__magic_name__					)

	UpperCamelCase     :List[str]          =	[int(2**32 * abs(sin(i + 1					)					)					) for i in range(64					)]

	# Starting states
	UpperCamelCase     :Union[str, Any]          =	0X67_45_23_01
	UpperCamelCase     :Union[str, Any]          =	0XEF_CD_AB_89
	UpperCamelCase     :List[str]          =	0X98_BA_DC_FE
	UpperCamelCase     :int          =	0X10_32_54_76

	UpperCamelCase     :int          =	[
	    7,
	    12,
	    17,
	    22,
	    7,
	    12,
	    17,
	    22,
	    7,
	    12,
	    17,
	    22,
	    7,
	    12,
	    17,
	    22,
	    5,
	    9,
	    14,
	    20,
	    5,
	    9,
	    14,
	    20,
	    5,
	    9,
	    14,
	    20,
	    5,
	    9,
	    14,
	    20,
	    4,
	    11,
	    16,
	    23,
	    4,
	    11,
	    16,
	    23,
	    4,
	    11,
	    16,
	    23,
	    4,
	    11,
	    16,
	    23,
	    6,
	    10,
	    15,
	    21,
	    6,
	    10,
	    15,
	    21,
	    6,
	    10,
	    15,
	    21,
	    6,
	    10,
	    15,
	    21,
	]

	# Process bit string in chunks, each with 16 32-char words
	for block_words in get_block_words(__magic_name__					):
		UpperCamelCase     :Optional[Any]          =	aa
		UpperCamelCase     :Any          =	ba
		UpperCamelCase     :Tuple          =	ca
		UpperCamelCase     :List[str]          =	da

		# Hash current chunk
		for i in range(64					):
			if i <= 15:
				# f = (b & c) | (not_32(b) & d)     # Alternate definition for f
				UpperCamelCase     :int          =	d ^ (b & (c ^ d))
				UpperCamelCase     :Optional[int]          =	i
			elif i <= 31:
				# f = (d & b) | (not_32(d) & c)     # Alternate definition for f
				UpperCamelCase     :str          =	c ^ (d & (b ^ c))
				UpperCamelCase     :Union[str, Any]          =	(5 * i + 1) % 16
			elif i <= 47:
				UpperCamelCase     :str          =	b ^ c ^ d
				UpperCamelCase     :Optional[int]          =	(3 * i + 5) % 16
			else:
				UpperCamelCase     :List[str]          =	c ^ (b | not_aa(__magic_name__					))
				UpperCamelCase     :int          =	(7 * i) % 16
			UpperCamelCase     :Dict          =	(f + a + added_consts[i] + block_words[g]) % 2**32
			UpperCamelCase     :Tuple          =	d
			UpperCamelCase     :str          =	c
			UpperCamelCase     :Tuple          =	b
			UpperCamelCase     :Optional[Any]          =	sum_aa(__magic_name__	,			left_rotate_aa(__magic_name__	,			shift_amounts[i]					)					)

		# Add hashed chunk to running total
		UpperCamelCase     :List[str]          =	sum_aa(__magic_name__	,			__magic_name__					)
		UpperCamelCase     :str          =	sum_aa(__magic_name__	,			__magic_name__					)
		UpperCamelCase     :int          =	sum_aa(__magic_name__	,			__magic_name__					)
		UpperCamelCase     :Optional[Any]          =	sum_aa(__magic_name__	,			__magic_name__					)

	UpperCamelCase     :Optional[Any]          =	reformat_hex(__magic_name__					) + reformat_hex(__magic_name__					) + reformat_hex(__magic_name__					) + reformat_hex(__magic_name__					)
	return digest


if __name__ == "__main__":
	import doctest

	doctest.testmod()
