




'''simple docstring'''





import tempfile
import unittest

import numpy as np
import torch

from diffusers import DiffusionPipeline
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor




class 				_snake_case	(			unittest.TestCase							):




			def      lowerCAmelCase__       (      self       ,							a__ = False							)       ->       Any:

										'''simple docstring'''




										snake_case_			      =   10

										if only_cross_attention:
																	snake_case_			      =   12
										else:
																	# when only cross attention is not set, the cross attention dim must be the same as the query dim
																	snake_case_			      =   query_dim

										return {
										    "query_dim": query_dim,
										    "cross_attention_dim": cross_attention_dim,
										    "heads": 2,
										    "dim_head": 4,
										    "added_kv_proj_dim": 6,
										    "norm_num_groups": 1,
										    "only_cross_attention": only_cross_attention,
										    "processor": AttnAddedKVProcessor(),
										}




			def      lowerCAmelCase__       (      self       ,							a__       ,							a__							)       ->       Dict:

										'''simple docstring'''




										snake_case_			      =   2

										snake_case_			      =   torch.rand(a__       ,							a__       ,							3       ,							2							)
										snake_case_			      =   torch.rand(a__       ,							4       ,							a__							)
										snake_case_			      =   None

										return {
										    "hidden_states": hidden_states,
										    "encoder_hidden_states": encoder_hidden_states,
										    "attention_mask": attention_mask,
										}




			def      lowerCAmelCase__       (      self							)       ->       List[Any]:

										'''simple docstring'''




										torch.manual_seed(0							)

										snake_case_			      =   self.get_constructor_arguments(only_cross_attention=a__							)
										snake_case_			      =   Attention(**a__							)

										self.assertTrue(attn.to_k is not None							)
										self.assertTrue(attn.to_v is not None							)

										snake_case_			      =   self.get_forward_arguments(
										    query_dim=constructor_args["query_dim"]       ,							added_kv_proj_dim=constructor_args["added_kv_proj_dim"]							)

										snake_case_			      =   attn(**a__							)

										# only self attention

										torch.manual_seed(0							)

										snake_case_			      =   self.get_constructor_arguments(only_cross_attention=a__							)
										snake_case_			      =   Attention(**a__							)

										self.assertTrue(attn.to_k is None							)
										self.assertTrue(attn.to_v is None							)

										snake_case_			      =   self.get_forward_arguments(
										    query_dim=constructor_args["query_dim"]       ,							added_kv_proj_dim=constructor_args["added_kv_proj_dim"]							)

										snake_case_			      =   attn(**a__							)

										self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all()							)





class 				_snake_case	(			unittest.TestCase							):




			def      lowerCAmelCase__       (      self							)       ->       Tuple:

										'''simple docstring'''




										snake_case_			      =   DiffusionPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-pipe"       ,							safety_checker=a__							)

										snake_case_			      =   pipe(
										    "foo"       ,							num_inference_steps=2       ,							generator=torch.Generator("cpu"							).manual_seed(0							)       ,							output_type="np"       ,							).images

										# the initial conversion succeeds
										snake_case_			      =   DiffusionPipeline.from_pretrained(
										    "hf-internal-testing/tiny-stable-diffusion-pipe"       ,							device_map="sequential"       ,							safety_checker=a__							)

										snake_case_			      =   pipe(
										    "foo"       ,							num_inference_steps=2       ,							generator=torch.Generator("cpu"							).manual_seed(0							)       ,							output_type="np"       ,							).images

										with tempfile.TemporaryDirectory() as tmpdir:
																	# save the converted model
																	pipe.save_pretrained(a__							)

																	# can also load the converted weights
																	snake_case_			      =   DiffusionPipeline.from_pretrained(a__       ,							device_map="sequential"       ,							safety_checker=a__							)

										snake_case_			      =   pipe(
										    "foo"       ,							num_inference_steps=2       ,							generator=torch.Generator("cpu"							).manual_seed(0							)       ,							output_type="np"       ,							).images

										self.assertTrue(np.allclose(a__       ,							a__       ,							atol=1e-5							)							)
										self.assertTrue(np.allclose(a__       ,							a__       ,							atol=1e-5							)							)










'''simple docstring'''





import heapq


def   UpperCamelCase_(			snake_case							: List[Any]							,       snake_case							: List[str]							,       snake_case							: Dict  ):

							'''simple docstring'''






							snake_case_			      =   [(0, start)]  # cost from start node,end node
							snake_case_			      =   set()
							while heap:
														((snake_case_)      ,   (snake_case_))			      =   heapq.heappop(snake_case  )
														if u in visited:
																					continue
														visited.add(snake_case  )
														if u == end:
																					return cost
														for v, c in graph[u]:
																					if v in visited:
																												continue
																					snake_case_			      =   cost + c
																					heapq.heappush(snake_case							,       (next_item, v)  )
							return -1


_SCREAMING_SNAKE_CASE							:     Dict   	=				{
    "A": [["B", 2], ["C", 5]],
    "B": [["A", 2], ["D", 3], ["E", 1], ["F", 1]],
    "C": [["A", 5], ["F", 3]],
    "D": [["B", 3]],
    "E": [["B", 4], ["F", 3]],
    "F": [["C", 3], ["E", 3]],
}
_SCREAMING_SNAKE_CASE							:     Optional[int]   	=				{
    "B": [["C", 1]],
    "C": [["D", 1]],
    "D": [["F", 1]],
    "E": [["B", 1], ["F", 3]],
    "F": [],
}
_SCREAMING_SNAKE_CASE							:     Any   	=				{
    "B": [["C", 1]],
    "C": [["D", 1]],
    "D": [["F", 1]],
    "E": [["B", 1], ["G", 2]],
    "F": [],
    "G": [["F", 1]],
}

_SCREAMING_SNAKE_CASE							:     Tuple   	=				dijkstra(G, "E", "C")
print(short_distance)  # E -- 3 --> F -- 3 --> C == 6

_SCREAMING_SNAKE_CASE							:     str   	=				dijkstra(Ga, "E", "F")
print(short_distance)  # E -- 3 --> F == 3

_SCREAMING_SNAKE_CASE							:     List[Any]   	=				dijkstra(Ga, "E", "F")
print(short_distance)  # E -- 2 --> G -- 1 --> F == 3

if __name__ == "__main__":
		import doctest

		doctest.testmod()




