


'''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 _lowerCamelCase	(						lowercase__		,					unittest.TestCase							):



		'''simple docstring'''







		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_     :      Any				   =    frozenset(
		    [
		        """num_inference_steps""",
		        """generator""",
		        """latents""",
		        """output_type""",
		        """return_dict""",
		        """callback""",
		        """callback_steps""",
		    ]							)







		@property
		def        __lowerCAmelCase		( self		:						List[str]      )  ->			int:
			__magic_name__	:   Optional[int]					     =					UNetaDModel.from_pretrained(
			    'diffusers/consistency-models-test'  ,					subfolder='test_unet'  ,					)
			return unet







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







		def        __lowerCAmelCase		( self		:						Tuple  ,					_A		:						Optional[int]=False      )  ->			Optional[int]:
			if class_cond:
				__magic_name__	:   List[Any]					     =					self.dummy_cond_unet
			else:
				__magic_name__	:   Optional[int]					     =					self.dummy_uncond_unet

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

			__magic_name__	:   Tuple					     =					{
			    'unet': unet,
			    'scheduler': scheduler,
			}

			return components







		def        __lowerCAmelCase		( self		:						List[Any]  ,					_A		:						Any  ,					_A		:						Optional[int]=0      )  ->			Optional[int]:
			if str(_A      ).startswith('mps'      ):
				__magic_name__	:   Any					     =					torch.manual_seed(_A      )
			else:
				__magic_name__	:   List[Any]					     =					torch.Generator(device=_A      ).manual_seed(_A      )

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

			return inputs







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

			__magic_name__	:   Optional[int]					     =					self.get_dummy_inputs(_A      )
			__magic_name__	:   List[Any]					     =					pipe(**_A      ).images
			assert image.shape == (1, 32, 32, 3)

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

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

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

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

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

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

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







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



		'''simple docstring'''







		def        __lowerCAmelCase		( self		:						Tuple      )  ->			Union[str, Any]:
			super().tearDown()
			gc.collect()
			torch.cuda.empty_cache()







		def        __lowerCAmelCase		( self		:						str  ,					_A		:						Optional[Any]=0  ,					_A		:						int=False  ,					_A		:						int="cpu"  ,					_A		:						Dict=torch.floataa  ,					_A		:						Optional[int]=(1, 3, 64, 64)      )  ->			Optional[Any]:
			__magic_name__	:   int					     =					torch.manual_seed(_A      )

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

			if get_fixed_latents:
				__magic_name__	:   str					     =					self.get_fixed_latents(seed=_A  ,					device=_A  ,					dtype=_A  ,					shape=_A      )
				__magic_name__	:   List[str]					     =					latents

			return inputs







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







		def        __lowerCAmelCase		( self		:						List[str]      )  ->			Any:
			__magic_name__	:   Optional[int]					     =					UNetaDModel.from_pretrained('diffusers/consistency_models'  ,					subfolder='diffusers_cd_imagenet64_l2'      )
			__magic_name__	:   Dict					     =					CMStochasticIterativeScheduler(
			    num_train_timesteps=40  ,					sigma_min=0.002  ,					sigma_max=80.0  ,					)
			__magic_name__	:   Union[str, Any]					     =					ConsistencyModelPipeline(unet=_A  ,					scheduler=_A      )
			pipe.to(torch_device=_A      )
			pipe.set_progress_bar_config(disable=_A      )

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

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

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

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

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

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

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

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

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

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

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

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



'''simple docstring'''

from __future__ import annotations

from sys import maxsize
from typing import Generic, TypeVar

lowerCAmelCase				:Any        =      TypeVar('''T''')





def lowerCamelCase      (      lowerCAmelCase    :     int		):





	"""simple docstring"""







	return (position - 1) // 2





def lowerCamelCase      (      lowerCAmelCase    :     int		):





	"""simple docstring"""







	return (2 * position) + 1





def lowerCamelCase      (      lowerCAmelCase    :     int		):





	"""simple docstring"""







	return (2 * position) + 2
class _lowerCamelCase	(						Generic[T]							):



		'''simple docstring'''







		def __init__( self		:						Union[str, Any]      )  ->			None:
			__magic_name__	:   list[tuple[T, int]]					     =					[]
			__magic_name__	:   dict[T, int]					     =					{}
			__magic_name__	:   int					     =					0







		def __len__( self		:						Optional[int]      )  ->			int:
			return self.elements







		def __repr__( self		:						Optional[Any]      )  ->			str:
			return str(self.heap      )







		def        __lowerCAmelCase		( self		:						List[Any]      )  ->			bool:
			# Check if the priority queue is empty
			return self.elements == 0







		def        __lowerCAmelCase		( self		:						Dict  ,					_A		:						T  ,					_A		:						int      )  ->			None:
			# Add an element with given priority to the queue
			self.heap.append((elem, weight)      )
			__magic_name__	:   List[str]					     =					self.elements
			self.elements += 1
			self._bubble_up(_A      )







		def        __lowerCAmelCase		( self		:						List[Any]      )  ->			T:
			# Remove and return the element with lowest weight (highest priority)
			if self.elements > 1:
				self._swap_nodes(0  ,					self.elements - 1      )
			__magic_name__     ,    __magic_name__	:   Optional[Any]					     =					self.heap.pop()
			del self.position_map[elem]
			self.elements -= 1
			if self.elements > 0:
				__magic_name__     ,    __magic_name__	:   Tuple					     =					self.heap[0]
				self._bubble_down(_A      )
			return elem







		def        __lowerCAmelCase		( self		:						Dict  ,					_A		:						T  ,					_A		:						int      )  ->			None:
			# Update the weight of the given key
			__magic_name__	:   Optional[Any]					     =					self.position_map[elem]
			__magic_name__	:   Any					     =					(elem, weight)
			if position > 0:
				__magic_name__	:   Optional[int]					     =					get_parent_position(_A      )
				__magic_name__     ,    __magic_name__	:   Any					     =					self.heap[parent_position]
				if parent_weight > weight:
					self._bubble_up(_A      )
				else:
					self._bubble_down(_A      )
			else:
				self._bubble_down(_A      )







		def        __lowerCAmelCase		( self		:						Optional[int]  ,					_A		:						T      )  ->			None:
			# Place a node at the proper position (upward movement) [to be used internally
			# only]
			__magic_name__	:   Union[str, Any]					     =					self.position_map[elem]
			if curr_pos == 0:
				return None
			__magic_name__	:   Optional[Any]					     =					get_parent_position(_A      )
			__magic_name__     ,    __magic_name__	:   Any					     =					self.heap[curr_pos]
			__magic_name__     ,    __magic_name__	:   int					     =					self.heap[parent_position]
			if parent_weight > weight:
				self._swap_nodes(_A  ,					_A      )
				return self._bubble_up(_A      )
			return None







		def        __lowerCAmelCase		( self		:						int  ,					_A		:						T      )  ->			None:
			# Place a node at the proper position (downward movement) [to be used
			# internally only]
			__magic_name__	:   Dict					     =					self.position_map[elem]
			__magic_name__     ,    __magic_name__	:   List[str]					     =					self.heap[curr_pos]
			__magic_name__	:   str					     =					get_child_left_position(_A      )
			__magic_name__	:   List[str]					     =					get_child_right_position(_A      )
			if child_left_position < self.elements and child_right_position < self.elements:
				__magic_name__     ,    __magic_name__	:   int					     =					self.heap[child_left_position]
				__magic_name__     ,    __magic_name__	:   int					     =					self.heap[child_right_position]
				if child_right_weight < child_left_weight and child_right_weight < weight:
					self._swap_nodes(_A  ,					_A      )
					return self._bubble_down(_A      )
			if child_left_position < self.elements:
				__magic_name__     ,    __magic_name__	:   Any					     =					self.heap[child_left_position]
				if child_left_weight < weight:
					self._swap_nodes(_A  ,					_A      )
					return self._bubble_down(_A      )
			else:
				return None
			if child_right_position < self.elements:
				__magic_name__     ,    __magic_name__	:   str					     =					self.heap[child_right_position]
				if child_right_weight < weight:
					self._swap_nodes(_A  ,					_A      )
					return self._bubble_down(_A      )
			return None







		def        __lowerCAmelCase		( self		:						int  ,					_A		:						int  ,					_A		:						int      )  ->			None:
			# Swap the nodes at the given positions
			__magic_name__	:   List[Any]					     =					self.heap[nodea_pos][0]
			__magic_name__	:   Union[str, Any]					     =					self.heap[nodea_pos][0]
			__magic_name__     ,    __magic_name__	:   Optional[Any]					     =					(
			    self.heap[nodea_pos],
			    self.heap[nodea_pos],
			)
			__magic_name__	:   Optional[Any]					     =					nodea_pos
			__magic_name__	:   Optional[Any]					     =					nodea_pos







class _lowerCamelCase	(						Generic[T]							):



		'''simple docstring'''







		def __init__( self		:						Union[str, Any]      )  ->			None:
			__magic_name__	:   dict[T, dict[T, int]]					     =					{}
			__magic_name__	:   int					     =					0







		def __repr__( self		:						str      )  ->			str:
			return str(self.connections      )







		def __len__( self		:						Union[str, Any]      )  ->			int:
			return self.nodes







		def        __lowerCAmelCase		( self		:						Dict  ,					_A		:						T      )  ->			None:
			# Add a node in the graph if it is not in the graph
			if node not in self.connections:
				__magic_name__	:   str					     =					{}
				self.nodes += 1







		def        __lowerCAmelCase		( self		:						List[str]  ,					_A		:						T  ,					_A		:						T  ,					_A		:						int      )  ->			None:
			# Add an edge between 2 nodes in the graph
			self.add_node(_A      )
			self.add_node(_A      )
			__magic_name__	:   Any					     =					weight
			__magic_name__	:   Dict					     =					weight







def lowerCamelCase      (      lowerCAmelCase    :     GraphUndirectedWeighted[T]		,		):





	"""simple docstring"""







	__magic_name__	:   dict[T, int]					     =					{node: maxsize for node in graph.connections}
	__magic_name__	:   dict[T, T | None]					     =					{node: None for node in graph.connections}

	__magic_name__	:   MinPriorityQueue[T]					     =					MinPriorityQueue()
	for node, weight in dist.items():
		priority_queue.push(lowerCAmelCase		,		lowerCAmelCase		)

	if priority_queue.is_empty():
		return dist, parent

	# initialization
	__magic_name__	:   Any					     =					priority_queue.extract_min()
	__magic_name__	:   List[str]					     =					0
	for neighbour in graph.connections[node]:
		if dist[neighbour] > dist[node] + graph.connections[node][neighbour]:
			__magic_name__	:   Any					     =					dist[node] + graph.connections[node][neighbour]
			priority_queue.update_key(lowerCAmelCase		,		dist[neighbour]		)
			__magic_name__	:   Tuple					     =					node

    # running prim's algorithm
	while not priority_queue.is_empty():
		__magic_name__	:   Tuple					     =					priority_queue.extract_min()
		for neighbour in graph.connections[node]:
			if dist[neighbour] > dist[node] + graph.connections[node][neighbour]:
				__magic_name__	:   int					     =					dist[node] + graph.connections[node][neighbour]
				priority_queue.update_key(lowerCAmelCase		,		dist[neighbour]		)
				__magic_name__	:   Union[str, Any]					     =					node
	return dist, parent