





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionAttendAndExcitePipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_numpy, skip_mps, slow
from diffusers.utils.testing_utils import require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


_a     :    Optional[int]         =						False
@skip_mps
class   a_ (			_A		,							_A		,							_A		,							unittest.TestCase      ):
			A__      :			Optional[Any]       	=  StableDiffusionAttendAndExcitePipeline
			A__      :			Optional[int]       	=  False
			A__      :			List[Any]       	=  TEXT_TO_IMAGE_PARAMS
			A__      :			Union[str, Any]       	=  TEXT_TO_IMAGE_BATCH_PARAMS.union({'token_indices'}      )
			A__      :			int       	=  TEXT_TO_IMAGE_IMAGE_PARAMS
			A__      :			Union[str, Any]       	=  TEXT_TO_IMAGE_IMAGE_PARAMS

			@classmethod
			def      lowerCAmelCase(   cls   :  Tuple						):






							"""simple docstring"""
							super().setUpClass()
							torch.use_deterministic_algorithms(a_						)

			@classmethod
			def      lowerCAmelCase(   cls   :  Tuple						):






							"""simple docstring"""
							super().tearDownClass()
							torch.use_deterministic_algorithms(a_						)

			def      lowerCAmelCase(   self   :  Optional[int]						):






							"""simple docstring"""
							torch.manual_seed(0						)
							snake_case   :					int    						=				UNetaDConditionModel(
							    block_out_channels=(32, 64)     , layers_per_block=1     , sample_size=32     , in_channels=4     , out_channels=4     , down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')     , up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')     , cross_attention_dim=32     , attention_head_dim=(2, 4)     , use_linear_projection=a_     , )
							snake_case   :					Union[str, Any]    						=				DDIMScheduler(
							    beta_start=0.0_0085     , beta_end=0.012     , beta_schedule='''scaled_linear'''     , clip_sample=a_     , set_alpha_to_one=a_     , )
							torch.manual_seed(0						)
							snake_case   :					Tuple    						=				AutoencoderKL(
							    block_out_channels=[32, 64]     , in_channels=3     , out_channels=3     , down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']     , up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']     , latent_channels=4     , sample_size=128     , )
							torch.manual_seed(0						)
							snake_case   :					int    						=				CLIPTextConfig(
							    bos_token_id=0     , eos_token_id=2     , hidden_size=32     , intermediate_size=37     , layer_norm_eps=1e-0_5     , num_attention_heads=4     , num_hidden_layers=5     , pad_token_id=1     , vocab_size=1_000     , hidden_act='''gelu'''     , projection_dim=512     , )
							snake_case   :					Tuple    						=				CLIPTextModel(a_						)
							snake_case   :					Union[str, Any]    						=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

							snake_case   :					Dict    						=				{
							    "unet": unet,
							    "scheduler": scheduler,
							    "vae": vae,
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "safety_checker": None,
							    "feature_extractor": None,
							}

							return components

			def      lowerCAmelCase(   self   :  str     , UpperCAmelCase__   :  Union[str, Any]     , UpperCAmelCase__   :  Optional[int]=0						):






							"""simple docstring"""
							if str(a_						).startswith('''mps'''						):
											snake_case   :					str    						=				torch.manual_seed(a_						)
							else:
											snake_case   :					int    						=				torch.Generator(device=a_						).manual_seed(a_						)
							snake_case   :					Any    						=				{
							    "prompt": "a cat and a frog",
							    "token_indices": [2, 5],
							    "generator": generator,
							    "num_inference_steps": 1,
							    "guidance_scale": 6.0,
							    "output_type": "numpy",
							    "max_iter_to_alter": 2,
							    "thresholds": {0: 0.7},
							}
							return inputs

			def      lowerCAmelCase(   self   :  Any						):






							"""simple docstring"""
							snake_case   :					str    						=				"cpu"

							snake_case   :					int    						=				self.get_dummy_components()
							snake_case   :					Union[str, Any]    						=				self.pipeline_class(**a_						)
							pipe.to(a_						)
							pipe.set_progress_bar_config(disable=a_						)

							snake_case   :					Optional[Any]    						=				self.get_dummy_inputs(a_						)
							snake_case   :					Optional[Any]    						=				pipe(**a_						).images
							snake_case   :					Tuple    						=				image[0, -3:, -3:, -1]

							self.assertEqual(image.shape     , (1, 64, 64, 3)						)
							snake_case   :					Union[str, Any]    						=				np.array(
							    [0.6390_5364, 0.6289_7307, 0.4859_9017, 0.513_3624, 0.555_0048, 0.4576_9516, 0.5032_6973, 0.502_3139, 0.4538_4496]						)
							snake_case   :					Tuple    						=				np.abs(image_slice.flatten() - expected_slice						).max()
							self.assertLessEqual(a_     , 1e-3						)

			def      lowerCAmelCase(   self   :  str						):






							"""simple docstring"""
							super().test_cpu_offload_forward_pass(expected_max_diff=5e-4						)

			def      lowerCAmelCase(   self   :  List[str]						):






							"""simple docstring"""
							# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
							self._test_inference_batch_consistent(batch_sizes=[1, 2]						)

			def      lowerCAmelCase(   self   :  Optional[Any]						):






							"""simple docstring"""
							self._test_inference_batch_single_identical(batch_size=2     , expected_max_diff=7e-4						)

			def      lowerCAmelCase(   self   :  Dict						):






							"""simple docstring"""
							super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3						)

			def      lowerCAmelCase(   self   :  List[str]						):






							"""simple docstring"""
							super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5e-4						)

			def      lowerCAmelCase(   self   :  Union[str, Any]						):






							"""simple docstring"""
							super().test_save_load_local(expected_max_difference=5e-4						)



			def      lowerCAmelCase(   self   :  Dict						):






							"""simple docstring"""
							super().test_save_load_optional_components(expected_max_difference=4e-4						)
@require_torch_gpu
@slow
class   a_ (			unittest.TestCase      ):

			@classmethod
			def      lowerCAmelCase(   cls   :  List[Any]						):






							"""simple docstring"""
							super().setUpClass()
							torch.use_deterministic_algorithms(a_						)

			@classmethod
			def      lowerCAmelCase(   cls   :  Any						):






							"""simple docstring"""
							super().tearDownClass()
							torch.use_deterministic_algorithms(a_						)

			def      lowerCAmelCase(   self   :  Dict						):






							"""simple docstring"""
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()



			def      lowerCAmelCase(   self   :  Dict						):






							"""simple docstring"""
							snake_case   :					List[Any]    						=				torch.manual_seed(51						)

							snake_case   :					List[Any]    						=				StableDiffusionAttendAndExcitePipeline.from_pretrained(
							    '''CompVis/stable-diffusion-v1-4'''     , safety_checker=a_     , torch_dtype=torch.floataa						)
							pipe.to('''cuda'''						)

							snake_case   :					List[str]    						=				"a painting of an elephant with glasses"
							snake_case   :					Dict    						=				[5, 7]

							snake_case   :					List[str]    						=				pipe(
							    prompt=a_     , token_indices=a_     , guidance_scale=7.5     , generator=a_     , num_inference_steps=5     , max_iter_to_alter=5     , output_type='''numpy'''     , ).images[0]

							snake_case   :					Any    						=				load_numpy(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/attend-and-excite/elephant_glasses.npy'''						)
							assert np.abs((expected_image - image).max()						) < 5e-1







'''simple docstring'''






from collections import defaultdict

from graphs.minimum_spanning_tree_prims import prisms_algorithm as mst






def   __A    (				):
		lowerCAmelCase						,		lowerCAmelCase  :     List[Any]       						= 9, 1_4  # noqa: F841
		lowerCAmelCase  :     int       						= [
		    [0, 1, 4],
		    [0, 7, 8],
		    [1, 2, 8],
		    [7, 8, 7],
		    [7, 6, 1],
		    [2, 8, 2],
		    [8, 6, 6],
		    [2, 3, 7],
		    [2, 5, 4],
		    [6, 5, 2],
		    [3, 5, 1_4],
		    [3, 4, 9],
		    [5, 4, 1_0],
		    [1, 7, 1_1],
		]

		lowerCAmelCase  :     Optional[Any]       						= defaultdict(a_						)
		for nodea, nodea, cost in edges:
				adjancency[nodea].append([nodea, cost]						)
				adjancency[nodea].append([nodea, cost]						)

		lowerCAmelCase  :     int       						= mst(a_						)

		lowerCAmelCase  :     Tuple       						= [
		    [7, 6, 1],
		    [2, 8, 2],
		    [6, 5, 2],
		    [0, 1, 4],
		    [2, 5, 4],
		    [2, 3, 7],
		    [0, 7, 8],
		    [3, 4, 9],
		]

		for answer in expected:
				lowerCAmelCase  :     List[str]       						= tuple(answer[:2]						)
				lowerCAmelCase  :     int       						= tuple(edge[::-1]						)
				assert edge in result or reverse in result



