



"""simple docstring"""

from __future__ import annotations

from typing import Any
class     lowerCAmelCase__						:


	'''simple docstring'''







	def __init__(      self				:       Optional[int]     ,       _SCREAMING_SNAKE_CASE				:       int    )						->	None:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     num_of_nodes
					SCREAMING_SNAKE_CASE						:							list[list[int]]							=     []
					SCREAMING_SNAKE_CASE						:							dict[int, int]							=     {}





	def 			_lowerCAmelCase     (      self				:       str     ,       _SCREAMING_SNAKE_CASE				:       int     ,       _SCREAMING_SNAKE_CASE				:       int     ,       _SCREAMING_SNAKE_CASE				:       int    )						->	None:





					"""simple docstring"""


					self.m_edges.append([u_node, v_node, weight]    )





	def 			_lowerCAmelCase     (      self				:       int     ,       _SCREAMING_SNAKE_CASE				:       int    )						->	int:





					"""simple docstring"""


					if self.m_component[u_node] == u_node:
									return u_node
					return self.find_component(self.m_component[u_node]    )





	def 			_lowerCAmelCase     (      self				:       Optional[Any]     ,       _SCREAMING_SNAKE_CASE				:       int    )						->	None:





					"""simple docstring"""


					if self.m_component[u_node] != u_node:
									for k in self.m_component:
													SCREAMING_SNAKE_CASE						:							Tuple							=     self.find_component(_SCREAMING_SNAKE_CASE    )





	def 			_lowerCAmelCase     (      self				:       List[str]     ,       _SCREAMING_SNAKE_CASE				:       list[int]     ,       _SCREAMING_SNAKE_CASE				:       int     ,       _SCREAMING_SNAKE_CASE				:       int    )						->	None:





					"""simple docstring"""


					if component_size[u_node] <= component_size[v_node]:
									SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     v_node
									component_size[v_node] += component_size[u_node]
									self.set_component(_SCREAMING_SNAKE_CASE    )

					elif component_size[u_node] >= component_size[v_node]:
									SCREAMING_SNAKE_CASE						:							List[str]							=     self.find_component(_SCREAMING_SNAKE_CASE    )
									component_size[u_node] += component_size[v_node]
									self.set_component(_SCREAMING_SNAKE_CASE    )







	def 			_lowerCAmelCase     (      self				:       int    )						->	None:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     []
					SCREAMING_SNAKE_CASE						:							Tuple							=     0

					SCREAMING_SNAKE_CASE						:							list[Any]							=     [-1] * self.m_num_of_nodes

					# A list of components (initialized to all of the nodes)
					for node in range(self.m_num_of_nodes    ):
									self.m_component.update({node: node}    )
									component_size.append(1    )

					SCREAMING_SNAKE_CASE						:							List[Any]							=     self.m_num_of_nodes

					while num_of_components > 1:
									for edge in self.m_edges:
													SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE						:							Optional[int]							=     edge

													SCREAMING_SNAKE_CASE						:							List[Any]							=     self.m_component[u]
													SCREAMING_SNAKE_CASE						:							Optional[int]							=     self.m_component[v]

													if u_component != v_component:

																	for component in (u_component, v_component):
																					if (
																					    minimum_weight_edge[component] == -1
																					    or minimum_weight_edge[component][2] > w
																					):
																									SCREAMING_SNAKE_CASE						:							int							=     [u, v, w]

									for edge in minimum_weight_edge:
													if isinstance(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE    ):
																	SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE				,						SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     edge

																	SCREAMING_SNAKE_CASE						:							int							=     self.m_component[u]
																	SCREAMING_SNAKE_CASE						:							int							=     self.m_component[v]

																	if u_component != v_component:
																					mst_weight += w
																					self.union(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE    )
																					print(f"""Added edge [{u} - {v}]\nAdded weight: {w}\n"""    )
																					num_of_components -= 1

									SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     [-1] * self.m_num_of_nodes
					print(f"""The total weight of the minimal spanning tree is: {mst_weight}"""    )




def  __snake_case							(	)     ->       None:






				'''simple docstring'''


if __name__ == "__main__":
							import doctest

							doctest.testmod()










"""simple docstring"""

import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()
@slow
@require_torch_gpu
class     lowerCAmelCase__						(	unittest.TestCase ):


	'''simple docstring'''







	def 			_lowerCAmelCase     (      self				:       List[Any]    )						->	str:





					"""simple docstring"""


					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()





	def 			_lowerCAmelCase     (      self				:       Any    )						->	str:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							int							=     StableDiffusionKDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'    )
					SCREAMING_SNAKE_CASE						:							Dict							=     sd_pipe.to(_SCREAMING_SNAKE_CASE    )
					sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					sd_pipe.set_scheduler('sample_euler'    )

					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     'A painting of a squirrel eating a burger'
					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     torch.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     sd_pipe([prompt]     ,       generator=_SCREAMING_SNAKE_CASE     ,       guidance_scale=9.0     ,       num_inference_steps=20     ,       output_type='np'    )

					SCREAMING_SNAKE_CASE						:							str							=     output.images

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

					assert image.shape == (1, 512, 512, 3)
					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     np.array([0.0_4_4_7, 0.0_4_9_2, 0.0_4_6_8, 0.0_4_0_8, 0.0_3_8_3, 0.0_4_0_8, 0.0_3_5_4, 0.0_3_8_0, 0.0_3_3_9]    )

					assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2





	def 			_lowerCAmelCase     (      self				:       Dict    )						->	Optional[int]:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     StableDiffusionKDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'    )
					SCREAMING_SNAKE_CASE						:							List[Any]							=     sd_pipe.to(_SCREAMING_SNAKE_CASE    )
					sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					sd_pipe.set_scheduler('sample_euler'    )

					SCREAMING_SNAKE_CASE						:							List[str]							=     'A painting of a squirrel eating a burger'
					SCREAMING_SNAKE_CASE						:							Optional[int]							=     torch.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							Dict							=     sd_pipe([prompt]     ,       generator=_SCREAMING_SNAKE_CASE     ,       guidance_scale=9.0     ,       num_inference_steps=20     ,       output_type='np'    )

					SCREAMING_SNAKE_CASE						:							Optional[int]							=     output.images

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

					assert image.shape == (1, 512, 512, 3)
					SCREAMING_SNAKE_CASE						:							List[Any]							=     np.array([0.1_2_3_7, 0.1_3_2_0, 0.1_4_3_8, 0.1_3_5_9, 0.1_3_9_0, 0.1_1_3_2, 0.1_2_7_7, 0.1_1_7_5, 0.1_1_1_2]    )

					assert np.abs(image_slice.flatten() - expected_slice    ).max() < 5E-1







	def 			_lowerCAmelCase     (      self				:       Dict    )						->	Optional[Any]:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE						:							Optional[int]							=     StableDiffusionKDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'    )
					SCREAMING_SNAKE_CASE						:							Union[str, Any]							=     sd_pipe.to(_SCREAMING_SNAKE_CASE    )
					sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE    )

					sd_pipe.set_scheduler('sample_dpmpp_2m'    )

					SCREAMING_SNAKE_CASE						:							Optional[int]							=     'A painting of a squirrel eating a burger'
					SCREAMING_SNAKE_CASE						:							str							=     torch.manual_seed(0    )
					SCREAMING_SNAKE_CASE						:							Optional[Any]							=     sd_pipe(
					    [prompt]     ,       generator=_SCREAMING_SNAKE_CASE     ,       guidance_scale=7.5     ,       num_inference_steps=15     ,       output_type='np'     ,       use_karras_sigmas=_SCREAMING_SNAKE_CASE     ,       )

					SCREAMING_SNAKE_CASE						:							int							=     output.images

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

					assert image.shape == (1, 512, 512, 3)
					SCREAMING_SNAKE_CASE						:							int							=     np.array(
					    [0.1_1_3_8_1_6_8_9, 0.1_2_1_1_2_9_2_1, 0.1_3_8_9_4_5_7, 0.1_2_5_4_9_6_0_6, 0.1_2_4_4_9_6_4, 0.1_0_8_3_1_5_1_7, 0.1_1_5_6_2_8_6_6, 0.1_0_8_6_7_8_1_6, 0.1_0_4_9_9_0_4_8]    )

					assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1E-2





