


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 	a			(						self       :						Union[str, Any]	,      _lowerCAmelCase       :						bool = False  )							->					int:



											'''simple docstring'''





											__magic_name__					=      1_0

											if only_cross_attention:
																	__magic_name__					=      1_2
											else:
																	# when only cross attention is not set, the cross attention dim must be the same as the query dim
																	__magic_name__					=      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 	a			(						self       :						Optional[Any]	,      _lowerCAmelCase       :						int	,      _lowerCAmelCase       :						str  )							->					str:



											'''simple docstring'''





											__magic_name__					=      2

											__magic_name__					=      torch.rand(_lowerCAmelCase	,      _lowerCAmelCase	,      3	,      2  )
											__magic_name__					=      torch.rand(_lowerCAmelCase	,      4	,      _lowerCAmelCase  )
											__magic_name__					=      None

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







					def 	a			(						self       :						List[Any]  )							->					Optional[int]:



											'''simple docstring'''





											torch.manual_seed(0  )

											__magic_name__					=      self.get_constructor_arguments(only_cross_attention=_lowerCAmelCase  )
											__magic_name__					=      Attention(**_lowerCAmelCase  )

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

											__magic_name__					=      self.get_forward_arguments(
											    query_dim=constructor_args['query_dim']	,      added_kv_proj_dim=constructor_args['added_kv_proj_dim']  )

											__magic_name__					=      attn(**_lowerCAmelCase  )

											# only self attention

											torch.manual_seed(0  )

											__magic_name__					=      self.get_constructor_arguments(only_cross_attention=_lowerCAmelCase  )
											__magic_name__					=      Attention(**_lowerCAmelCase  )

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

											__magic_name__					=      self.get_forward_arguments(
											    query_dim=constructor_args['query_dim']	,      added_kv_proj_dim=constructor_args['added_kv_proj_dim']  )

											__magic_name__					=      attn(**_lowerCAmelCase  )

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







class      __snake_case				(   unittest.TestCase):
					def 	a			(						self       :						List[Any]  )							->					List[Any]:



											'''simple docstring'''





											__magic_name__					=      DiffusionPipeline.from_pretrained('hf-internal-testing/tiny-stable-diffusion-pipe'	,      safety_checker=_lowerCAmelCase  )

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

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

											__magic_name__					=      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(_lowerCAmelCase  )

																	# can also load the converted weights
																	__magic_name__					=      DiffusionPipeline.from_pretrained(_lowerCAmelCase	,      device_map='sequential'	,      safety_checker=_lowerCAmelCase  )

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

											self.assertTrue(np.allclose(_lowerCAmelCase	,      _lowerCAmelCase	,      atol=1E-5  )  )
											self.assertTrue(np.allclose(_lowerCAmelCase	,      _lowerCAmelCase	,      atol=1E-5  )  )








def        __SCREAMING_SNAKE_CASE	(				__lowercase:		int):



  '''simple docstring'''






  if not isinstance(__lowercase  ,				__lowercase):
    __lowerCAmelCase								=					F'''Input value of [number={number}] must be an integer'''
    raise TypeError(__lowercase)
  if number < 0 or number % 2 == 1:
    return False
  __lowerCAmelCase								=					int(number**0.5)
  return number == number_sqrt * (number_sqrt + 1)


if __name__ == "__main__":
      import doctest

      doctest.testmod()
