"""simple docstring"""







import gc
import random
import unittest

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

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


def      __lowerCAmelCase(		__UpperCAmelCase			):






	"""simple docstring"""





	_lowercase :    Dict							=							[tensor.shape for tensor in tensor_list]
	return all(shape == shapes[0] for shape in shapes[1:]			)



class  _lowerCamelCase			(__lowerCamelCase						,							__lowerCamelCase						,							__lowerCamelCase						,							unittest.TestCase     ):
				_snake_case     =  StableDiffusionLatentUpscalePipeline
				_snake_case     =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
				    "height",
				    "width",
				    "cross_attention_kwargs",
				    "negative_prompt_embeds",
				    "prompt_embeds",
				}
				_snake_case     =  PipelineTesterMixin.required_optional_params - {"num_images_per_prompt"}
				_snake_case     =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
				_snake_case     =  frozenset(
				    []     )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
				_snake_case     =  frozenset([]     )

				_snake_case     =  True







				@property
				def 							__UpperCAmelCase					( self    :       Optional[Any]						):

					"""simple docstring"""






					_lowercase :    str							=							1
					_lowercase :    Any							=							4
					_lowercase :    Optional[int]							=							(1_6, 1_6)

					_lowercase :    Dict							=							floats_tensor((batch_size, num_channels) + sizes			,					rng=random.Random(0						)						).to(lowerCamelCase_						)
					return image







				def 							__UpperCAmelCase					( self    :       Tuple						):

					"""simple docstring"""






					torch.manual_seed(0						)
					_lowercase :    Tuple							=							UNetaDConditionModel(
					    act_fn='gelu'			,					attention_head_dim=8			,					norm_num_groups=lowerCamelCase_			,					block_out_channels=[3_2, 3_2, 6_4, 6_4]			,					time_cond_proj_dim=1_6_0			,					conv_in_kernel=1			,					conv_out_kernel=1			,					cross_attention_dim=3_2			,					down_block_types=(
					        'KDownBlock2D',
					        'KCrossAttnDownBlock2D',
					        'KCrossAttnDownBlock2D',
					        'KCrossAttnDownBlock2D',
					    )			,					in_channels=8			,					mid_block_type=lowerCamelCase_			,					only_cross_attention=lowerCamelCase_			,					out_channels=5			,					resnet_time_scale_shift='scale_shift'			,					time_embedding_type='fourier'			,					timestep_post_act='gelu'			,					up_block_types=('KCrossAttnUpBlock2D', 'KCrossAttnUpBlock2D', 'KCrossAttnUpBlock2D', 'KUpBlock2D')			,					)
					_lowercase :    Optional[Any]							=							AutoencoderKL(
					    block_out_channels=[3_2, 3_2, 6_4, 6_4]			,					in_channels=3			,					out_channels=3			,					down_block_types=[
					        'DownEncoderBlock2D',
					        'DownEncoderBlock2D',
					        'DownEncoderBlock2D',
					        'DownEncoderBlock2D',
					    ]			,					up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D']			,					latent_channels=4			,					)
					_lowercase :    List[Any]							=							EulerDiscreteScheduler(prediction_type='sample'						)
					_lowercase :    str							=							CLIPTextConfig(
					    bos_token_id=0			,					eos_token_id=2			,					hidden_size=3_2			,					intermediate_size=3_7			,					layer_norm_eps=1E-05			,					num_attention_heads=4			,					num_hidden_layers=5			,					pad_token_id=1			,					vocab_size=1_0_0_0			,					hidden_act='quick_gelu'			,					projection_dim=5_1_2			,					)
					_lowercase :    int							=							CLIPTextModel(lowerCamelCase_						)
					_lowercase :    Tuple							=							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

					_lowercase :    int							=							{
					    'unet': model.eval(),
					    'vae': vae.eval(),
					    'scheduler': scheduler,
					    'text_encoder': text_encoder,
					    'tokenizer': tokenizer,
					}

					return components







				def 							__UpperCAmelCase					( self    :       Any			,					lowerCamelCase_    :       str			,					lowerCamelCase_    :       List[str]=0						):

					"""simple docstring"""






					if str(lowerCamelCase_						).startswith('mps'						):
						_lowercase :    Optional[int]							=							torch.manual_seed(lowerCamelCase_						)
					else:
						_lowercase :    Dict							=							torch.Generator(device=lowerCamelCase_						).manual_seed(lowerCamelCase_						)
					_lowercase :    Optional[int]							=							{
					    'prompt': 'A painting of a squirrel eating a burger',
					    'image': self.dummy_image.cpu(),
					    'generator': generator,
					    'num_inference_steps': 2,
					    'output_type': 'numpy',
					}
					return inputs







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

					"""simple docstring"""






					_lowercase :    int							=							'cpu'

					_lowercase :    Dict							=							self.get_dummy_components()
					_lowercase :    Optional[int]							=							self.pipeline_class(**lowerCamelCase_						)
					pipe.to(lowerCamelCase_						)
					pipe.set_progress_bar_config(disable=lowerCamelCase_						)

					_lowercase :    List[Any]							=							self.get_dummy_inputs(lowerCamelCase_						)
					_lowercase :    Optional[Any]							=							pipe(**lowerCamelCase_						).images
					_lowercase :    List[str]							=							image[0, -3:, -3:, -1]

					self.assertEqual(image.shape			,					(1, 2_5_6, 2_5_6, 3)						)
					_lowercase :    List[Any]							=							np.array(
					    [0.4722_2412, 0.4192_1633, 0.4471_7434, 0.4687_4192, 0.4258_8258, 0.4615_0726, 0.467_7534, 0.4558_3832, 0.4857_9055]						)
					_lowercase :    str							=							np.abs(image_slice.flatten() - expected_slice						).max()
					self.assertLessEqual(lowerCamelCase_			,					1E-3						)







				def 							__UpperCAmelCase					( self    :       Optional[int]						):

					"""simple docstring"""






					super().test_attention_slicing_forward_pass(expected_max_diff=7E-3						)







				def 							__UpperCAmelCase					( self    :       Optional[Any]						):

					"""simple docstring"""






					super().test_cpu_offload_forward_pass(expected_max_diff=3E-3						)







				def 							__UpperCAmelCase					( self    :       Any						):

					"""simple docstring"""






					super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3						)







				def 							__UpperCAmelCase					( self    :       int						):

					"""simple docstring"""






					super().test_inference_batch_single_identical(expected_max_diff=7E-3						)







				def 							__UpperCAmelCase					( self    :       Any						):

					"""simple docstring"""






					super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3E-3						)







				def 							__UpperCAmelCase					( self    :       int						):

					"""simple docstring"""






					super().test_save_load_local(expected_max_difference=3E-3						)







				def 							__UpperCAmelCase					( self    :       str						):

					"""simple docstring"""






					super().test_save_load_optional_components(expected_max_difference=3E-3						)







				def 							__UpperCAmelCase					( self    :       str						):

					"""simple docstring"""






					_lowercase :    Dict							=							[
					    'DDIMScheduler',
					    'DDPMScheduler',
					    'PNDMScheduler',
					    'HeunDiscreteScheduler',
					    'EulerAncestralDiscreteScheduler',
					    'KDPM2DiscreteScheduler',
					    'KDPM2AncestralDiscreteScheduler',
					    'DPMSolverSDEScheduler',
					]
					_lowercase :    Optional[Any]							=							self.get_dummy_components()
					_lowercase :    Dict							=							self.pipeline_class(**lowerCamelCase_						)

					# make sure that PNDM does not need warm-up
					pipe.scheduler.register_to_config(skip_prk_steps=lowerCamelCase_						)

					pipe.to(lowerCamelCase_						)
					pipe.set_progress_bar_config(disable=lowerCamelCase_						)
					_lowercase :    Optional[Any]							=							self.get_dummy_inputs(lowerCamelCase_						)
					_lowercase :    Optional[int]							=							2

					_lowercase :    List[str]							=							[]
					for scheduler_enum in KarrasDiffusionSchedulers:
						if scheduler_enum.name in skip_schedulers:
							# no sigma schedulers are not supported
							# no schedulers
							continue

						_lowercase :    Dict							=							getattr(lowerCamelCase_			,					scheduler_enum.name						)
						_lowercase :    int							=							scheduler_cls.from_config(pipe.scheduler.config						)
						_lowercase :    Optional[Any]							=							pipe(**lowerCamelCase_						)[0]
						outputs.append(lowerCamelCase_						)

					assert check_same_shape(lowerCamelCase_						)



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







				def 							__UpperCAmelCase					( self    :       Optional[Any]						):

					"""simple docstring"""






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







				def 							__UpperCAmelCase					( self    :       Any						):

					"""simple docstring"""






					_lowercase :    Optional[int]							=							torch.manual_seed(3_3						)

					_lowercase :    Tuple							=							StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'			,					torch_dtype=torch.floataa						)
					pipe.to('cuda'						)

					_lowercase :    Any							=							StableDiffusionLatentUpscalePipeline.from_pretrained(
					    'stabilityai/sd-x2-latent-upscaler'			,					torch_dtype=torch.floataa						)
					upscaler.to('cuda'						)

					_lowercase :    Dict							=							'a photo of an astronaut high resolution, unreal engine, ultra realistic'

					_lowercase :    Tuple							=							pipe(lowerCamelCase_			,					generator=lowerCamelCase_			,					output_type='latent'						).images

					_lowercase :    str							=							upscaler(
					    prompt=lowerCamelCase_			,					image=lowerCamelCase_			,					num_inference_steps=2_0			,					guidance_scale=0			,					generator=lowerCamelCase_			,					output_type='np'			,					).images[0]

					_lowercase :    Union[str, Any]							=							load_numpy(
					    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy'						)
					assert np.abs((expected_image - image).mean()						) < 5E-2







				def 							__UpperCAmelCase					( self    :       int						):

					"""simple docstring"""






					_lowercase :    List[Any]							=							torch.manual_seed(3_3						)

					_lowercase :    str							=							StableDiffusionLatentUpscalePipeline.from_pretrained(
					    'stabilityai/sd-x2-latent-upscaler'			,					torch_dtype=torch.floataa						)
					upscaler.to('cuda'						)

					_lowercase :    Optional[int]							=							'the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas'

					_lowercase :    int							=							load_image(
					    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png'						)

					_lowercase :    List[str]							=							upscaler(
					    prompt=lowerCamelCase_			,					image=lowerCamelCase_			,					num_inference_steps=2_0			,					guidance_scale=0			,					generator=lowerCamelCase_			,					output_type='np'			,					).images[0]

					_lowercase :    Optional[Any]							=							load_numpy(
					    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy'						)
					assert np.abs((expected_image - image).max()						) < 5E-2




"""simple docstring"""







from __future__ import annotations

SCREAMING_SNAKE_CASE									=						[
    [-1, 0],  # left
    [0, -1],  # down
    [1, 0],  # right
    [0, 1],  # up
]


def      __lowerCAmelCase(		__UpperCAmelCase     ,__UpperCAmelCase     ,__UpperCAmelCase     ,__UpperCAmelCase     ,__UpperCAmelCase     ,):






	"""simple docstring"""





	_lowercase :    List[str]							=							[
	    [0 for col in range(len(grid[0]			)			)] for row in range(len(__UpperCAmelCase			)			)
	]  # the reference grid
	_lowercase :    Optional[Any]							=							1
	_lowercase :    Any							=							[
	    [0 for col in range(len(grid[0]			)			)] for row in range(len(__UpperCAmelCase			)			)
	]  # the action grid

	_lowercase :    List[Any]							=							init[0]
	_lowercase :    Optional[Any]							=							init[1]
	_lowercase :    Union[str, Any]							=							0
	_lowercase :    Union[str, Any]							=							g + heuristic[x][y]  # cost from starting cell to destination cell
	_lowercase :    Dict							=							[[f, g, x, y]]

	_lowercase :    Tuple							=							False  # flag that is set when search is complete
	_lowercase :    Tuple							=							False  # flag set if we can't find expand

	while not found and not resign:
		if len(__UpperCAmelCase			) == 0:
			raise ValueError('Algorithm is unable to find solution'			)
		else:  # to choose the least costliest action so as to move closer to the goal
			cell.sort()
			cell.reverse()
			_lowercase :    List[Any]							=							cell.pop()
			_lowercase :    Optional[Any]							=							next_cell[2]
			_lowercase :    List[str]							=							next_cell[3]
			_lowercase :    Optional[int]							=							next_cell[1]

			if x == goal[0] and y == goal[1]:
				_lowercase :    Optional[Any]							=							True
			else:
				for i in range(len(__UpperCAmelCase			)			):  # to try out different valid actions
					_lowercase :    Any							=							x + DIRECTIONS[i][0]
					_lowercase :    Union[str, Any]							=							y + DIRECTIONS[i][1]
					if xa >= 0 and xa < len(__UpperCAmelCase			) and ya >= 0 and ya < len(grid[0]			):
						if closed[xa][ya] == 0 and grid[xa][ya] == 0:
							_lowercase :    int							=							g + cost
							_lowercase :    List[str]							=							ga + heuristic[xa][ya]
							cell.append([fa, ga, xa, ya]			)
							_lowercase :    List[Any]							=							1
							_lowercase :    Dict							=							i
	_lowercase :    Union[str, Any]							=							[]
	_lowercase :    Optional[int]							=							goal[0]
	_lowercase :    Any							=							goal[1]
	invpath.append([x, y]			)  # we get the reverse path from here
	while x != init[0] or y != init[1]:
		_lowercase :    str							=							x - DIRECTIONS[action[x][y]][0]
		_lowercase :    Any							=							y - DIRECTIONS[action[x][y]][1]
		_lowercase :    Dict							=							xa
		_lowercase :    Tuple							=							ya
		invpath.append([x, y]			)

	_lowercase :    List[str]							=							[]
	for i in range(len(__UpperCAmelCase			)			):
		path.append(invpath[len(__UpperCAmelCase			) - 1 - i]			)
	return path, action


if __name__ == "__main__":
						SCREAMING_SNAKE_CASE									=						[
						    [0, 1, 0, 0, 0, 0],
						    [0, 1, 0, 0, 0, 0],  # 0 are free path whereas 1's are obstacles
						    [0, 1, 0, 0, 0, 0],
						    [0, 1, 0, 0, 1, 0],
						    [0, 0, 0, 0, 1, 0],
						]

						SCREAMING_SNAKE_CASE									=						[0, 0]
						# all coordinates are given in format [y,x]
						SCREAMING_SNAKE_CASE									=						[len(grid) - 1, len(grid[0]) - 1]
						SCREAMING_SNAKE_CASE									=						1

						# the cost map which pushes the path closer to the goal
						SCREAMING_SNAKE_CASE									=						[[0 for row in range(len(grid[0]))] for col in range(len(grid))]
						for i in range(len(grid)):
												for j in range(len(grid[0])):
																		SCREAMING_SNAKE_CASE									=						abs(i - goal[0]) + abs(j - goal[1])
																		if grid[i][j] == 1:
																								# added extra penalty in the heuristic map
																								SCREAMING_SNAKE_CASE									=						99

						SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE									=						search(grid, init, goal, cost, heuristic)

						print('ACTION MAP')
						for i in range(len(action)):
												print(action[i])

						for i in range(len(path)):
												print(path[i])



