import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

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


enable_full_determinism()
@skip_mps
class 							a_					(		A_			,	A_			,	A_			,	unittest.TestCase		):






 '''simple docstring'''






 UpperCAmelCase_      					=   StableDiffusionModelEditingPipeline
 UpperCAmelCase_      					=   TEXT_TO_IMAGE_PARAMS
 UpperCAmelCase_      					=   TEXT_TO_IMAGE_BATCH_PARAMS
 UpperCAmelCase_      					=   TEXT_TO_IMAGE_IMAGE_PARAMS
 UpperCAmelCase_      					=   TEXT_TO_IMAGE_IMAGE_PARAMS




 def 						__snake_case		(     self   :    List[str]):







        '''simple docstring'''



        torch.manual_seed(0)
        lowerCAmelCase__	      =   UNetaDConditionModel(
            block_out_channels=(32, 64)     ,       layers_per_block=2     ,       sample_size=32     ,       in_channels=4     ,       out_channels=4     ,       down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')     ,       up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')     ,       cross_attention_dim=32     ,       )
        lowerCAmelCase__	      =   DDIMScheduler()
        torch.manual_seed(0)
        lowerCAmelCase__	      =   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     ,       )
        torch.manual_seed(0)
        lowerCAmelCase__	      =   CLIPTextConfig(
            bos_token_id=0     ,       eos_token_id=2     ,       hidden_size=32     ,       intermediate_size=37     ,       layer_norm_eps=1e-05     ,       num_attention_heads=4     ,       num_hidden_layers=5     ,       pad_token_id=1     ,       vocab_size=1_000     ,       )
        lowerCAmelCase__	      =   CLIPTextModel(lowercase__)
        lowerCAmelCase__	      =   CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

        lowerCAmelCase__	      =   {
            'unet': unet,
            'scheduler': scheduler,
            'vae': vae,
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'safety_checker': None,
            'feature_extractor': None,
        }
        return components




 def 						__snake_case		(     self   :    Optional[Any]     ,       lowercase__   :    int     ,       lowercase__   :    str=0):







        '''simple docstring'''



        lowerCAmelCase__	      =   torch.manual_seed(lowercase__)
        lowerCAmelCase__	      =   {
            'prompt': 'A field of roses',
            'generator': generator,
            # Setting height and width to None to prevent OOMs on CPU.
            'height': None,
            'width': None,
            'num_inference_steps': 2,
            'guidance_scale': 6.0,
            'output_type': 'numpy',
        }
        return inputs




 def 						__snake_case		(     self   :    List[Any]):







        '''simple docstring'''



        lowerCAmelCase__	      =   'cpu'  # ensure determinism for the device-dependent torch.Generator
        lowerCAmelCase__	      =   self.get_dummy_components()
        lowerCAmelCase__	      =   StableDiffusionModelEditingPipeline(**lowercase__)
        lowerCAmelCase__	      =   sd_pipe.to(lowercase__)
        sd_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__	      =   self.get_dummy_inputs(lowercase__)
        lowerCAmelCase__	      =   sd_pipe(**lowercase__).images
        lowerCAmelCase__	      =   image[0, -3:, -3:, -1]
        assert image.shape == (1, 64, 64, 3)

        lowerCAmelCase__	      =   np.array([0.4_755, 0.5_132, 0.4_976, 0.3_904, 0.3_554, 0.4_765, 0.5_139, 0.5_158, 0.4_889])

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




 def 						__snake_case		(     self   :    Optional[int]):







        '''simple docstring'''



        lowerCAmelCase__	      =   'cpu'  # ensure determinism for the device-dependent torch.Generator
        lowerCAmelCase__	      =   self.get_dummy_components()
        lowerCAmelCase__	      =   StableDiffusionModelEditingPipeline(**lowercase__)
        lowerCAmelCase__	      =   sd_pipe.to(lowercase__)
        sd_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__	      =   self.get_dummy_inputs(lowercase__)
        lowerCAmelCase__	      =   'french fries'
        lowerCAmelCase__	      =   sd_pipe(**lowercase__     ,       negative_prompt=lowercase__)
        lowerCAmelCase__	      =   output.images
        lowerCAmelCase__	      =   image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)

        lowerCAmelCase__	      =   np.array([0.4_992, 0.5_101, 0.5_004, 0.3_949, 0.3_604, 0.4_735, 0.5_216, 0.5_204, 0.4_913])

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




 def 						__snake_case		(     self   :    Tuple):







        '''simple docstring'''



        lowerCAmelCase__	      =   'cpu'  # ensure determinism for the device-dependent torch.Generator
        lowerCAmelCase__	      =   self.get_dummy_components()
        lowerCAmelCase__	      =   EulerAncestralDiscreteScheduler(
            beta_start=0.00_085     ,       beta_end=0.012     ,       beta_schedule='scaled_linear')
        lowerCAmelCase__	      =   StableDiffusionModelEditingPipeline(**lowercase__)
        lowerCAmelCase__	      =   sd_pipe.to(lowercase__)
        sd_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__	      =   self.get_dummy_inputs(lowercase__)
        lowerCAmelCase__	      =   sd_pipe(**lowercase__).images
        lowerCAmelCase__	      =   image[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)

        lowerCAmelCase__	      =   np.array([0.4_747, 0.5_372, 0.4_779, 0.4_982, 0.5_543, 0.4_816, 0.5_238, 0.4_904, 0.5_027])

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




 def 						__snake_case		(     self   :    int):







        '''simple docstring'''



        lowerCAmelCase__	      =   'cpu'  # ensure determinism for the device-dependent torch.Generator
        lowerCAmelCase__	      =   self.get_dummy_components()
        lowerCAmelCase__	      =   PNDMScheduler()
        lowerCAmelCase__	      =   StableDiffusionModelEditingPipeline(**lowercase__)
        lowerCAmelCase__	      =   sd_pipe.to(lowercase__)
        sd_pipe.set_progress_bar_config(disable=lowercase__)

        lowerCAmelCase__	      =   self.get_dummy_inputs(lowercase__)
        # the pipeline does not expect pndm so test if it raises error.
        with self.assertRaises(lowercase__):
               lowerCAmelCase__	      =   sd_pipe(**lowercase__).images




 def 						__snake_case		(     self   :    Tuple):







        '''simple docstring'''



        super().test_inference_batch_single_identical(expected_max_diff=5e-3)




 def 						__snake_case		(     self   :    List[Any]):







        '''simple docstring'''



        super().test_attention_slicing_forward_pass(expected_max_diff=5e-3)






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






 '''simple docstring'''






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







        '''simple docstring'''



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




 def 						__snake_case		(     self   :    Union[str, Any]     ,       lowercase__   :    Union[str, Any]=0):







        '''simple docstring'''



        lowerCAmelCase__	      =   torch.manual_seed(lowercase__)
        lowerCAmelCase__	      =   {
            'prompt': 'A field of roses',
            'generator': generator,
            'num_inference_steps': 3,
            'guidance_scale': 7.5,
            'output_type': 'numpy',
        }
        return inputs




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







        '''simple docstring'''



        lowerCAmelCase__	      =   'CompVis/stable-diffusion-v1-4'
        lowerCAmelCase__	      =   StableDiffusionModelEditingPipeline.from_pretrained(lowercase__     ,       safety_checker=lowercase__)
        pipe.to(lowercase__)
        pipe.set_progress_bar_config(disable=lowercase__)
        pipe.enable_attention_slicing()

        lowerCAmelCase__	      =   self.get_inputs()
        lowerCAmelCase__	      =   pipe(**lowercase__).images
        lowerCAmelCase__	      =   image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 512, 512, 3)

        lowerCAmelCase__	      =   np.array(
            [0.6_749_496, 0.6_386_453, 0.51_443_267, 0.66_094_905, 0.61_921_215, 0.5_491_332, 0.5_744_417, 0.58_075_106, 0.5_174_658])

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

        # make sure image changes after editing
        pipe.edit_model('A pack of roses'     ,       'A pack of blue roses')

        lowerCAmelCase__	      =   pipe(**lowercase__).images
        lowerCAmelCase__	      =   image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 512, 512, 3)

        assert np.abs(expected_slice - image_slice).max() > 1e-1




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







        '''simple docstring'''



        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        lowerCAmelCase__	      =   'CompVis/stable-diffusion-v1-4'
        lowerCAmelCase__	      =   DDIMScheduler.from_pretrained(lowercase__     ,       subfolder='scheduler')
        lowerCAmelCase__	      =   StableDiffusionModelEditingPipeline.from_pretrained(
            lowercase__     ,       scheduler=lowercase__     ,       safety_checker=lowercase__)
        lowerCAmelCase__	      =   pipe.to(lowercase__)
        pipe.set_progress_bar_config(disable=lowercase__)
        pipe.enable_attention_slicing(1)
        pipe.enable_sequential_cpu_offload()

        lowerCAmelCase__	      =   self.get_inputs()
        lowerCAmelCase__	      =   pipe(**lowercase__)

        lowerCAmelCase__	      =   torch.cuda.max_memory_allocated()
        # make sure that less than 4.4 GB is allocated
        assert mem_bytes < 4.4 * 10**9





def    UpperCAmelCase					(  lowercase__						:					List[Any]				,					lowercase__						:					Optional[int]	):



			'''simple docstring'''





			print("""\nVertex Distance"""	)
			for i in range(lowercase__	):
						if dist[i] != float("""inf"""	):
									print(lowercase__				,					"""\t"""				,					int(dist[i]	)				,					end="""\t"""	)
						else:
									print(lowercase__				,					"""\t"""				,					"""INF"""				,					end="""\t"""	)
						print()




def    UpperCAmelCase					(  lowercase__						:					Optional[int]				,					lowercase__						:					Any				,					lowercase__						:					Tuple	):



			'''simple docstring'''





			a__					=  float("""inf"""	)
			a__					=  -1
			for i in range(lowercase__	):
						if (not vset[i]) and mdist[i] < min_val:
									a__					=  i
									a__					=  mdist[i]
			return min_ind




def    UpperCAmelCase					(  lowercase__						:					Optional[Any]				,					lowercase__						:					Tuple				,					lowercase__						:					List[Any]	):



			'''simple docstring'''





			a__					=  [float("""inf"""	) for _ in range(lowercase__	)]
			a__					=  [False for _ in range(lowercase__	)]
			a__					=  0.0

			for _ in range(v - 1	):
						a__					=  min_dist(lowercase__				,					lowercase__				,					lowercase__	)
						a__					=  True

						for i in range(lowercase__	):
									if (
									    (not vset[i])
									    and graph[u][i] != float("""inf"""	)
									    and mdist[u] + graph[u][i] < mdist[i]
									):
												a__					=  mdist[u] + graph[u][i]

			print_dist(lowercase__				,					lowercase__	)


if __name__ == "__main__":
	_lowercase				:					Any     					=int(input("""Enter number of vertices: """).strip())
	_lowercase				:					List[str]     					=int(input("""Enter number of edges: """).strip())

	_lowercase				:					List[str]     					=[[float("""inf""") for i in range(V)] for j in range(V)]

	for i in range(V):
		_lowercase				:					Dict     					=0.0

	for i in range(E):
		print("""\nEdge """, i + 1)
		_lowercase				:					Dict     					=int(input("""Enter source:""").strip())
		_lowercase				:					Optional[Any]     					=int(input("""Enter destination:""").strip())
		_lowercase				:					List[Any]     					=float(input("""Enter weight:""").strip())
		_lowercase				:					List[str]     					=weight

	_lowercase				:					List[Any]     					=int(input("""\nEnter shortest path source:""").strip())
	dijkstra(graph, V, gsrc)

