
"""simple docstring"""

import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()






class 	_lowercase       (    __UpperCAmelCase    ,	unittest.TestCase		):
	lowercase_				=     PriorTransformer
	lowercase_				=     'hidden_states'




	@property
	def       _UpperCamelCase   (			self )		->							List[str]:
			lowerCamelCase					:      Any          =   4
			lowerCamelCase					:      Dict          =   8
			lowerCamelCase					:      Tuple          =   7

			lowerCamelCase					:      List[str]          =   floats_tensor((batch_size, embedding_dim) ).to(UpperCAmelCase_ )

			lowerCamelCase					:      List[str]          =   floats_tensor((batch_size, embedding_dim) ).to(UpperCAmelCase_ )
			lowerCamelCase					:      int          =   floats_tensor((batch_size, num_embeddings, embedding_dim) ).to(UpperCAmelCase_ )

			return {
			    "hidden_states": hidden_states,
			    "timestep": 2,
			    "proj_embedding": proj_embedding,
			    "encoder_hidden_states": encoder_hidden_states,
			}




	def       _UpperCamelCase   (			self		,			UpperCAmelCase_=0 )		->							str:
			torch.manual_seed(UpperCAmelCase_ )
			lowerCamelCase					:      Tuple          =   4
			lowerCamelCase					:      Optional[int]          =   8
			lowerCamelCase					:      Union[str, Any]          =   7

			lowerCamelCase					:      Optional[int]          =   torch.randn((batch_size, embedding_dim) ).to(UpperCAmelCase_ )

			lowerCamelCase					:      Optional[Any]          =   torch.randn((batch_size, embedding_dim) ).to(UpperCAmelCase_ )
			lowerCamelCase					:      Dict          =   torch.randn((batch_size, num_embeddings, embedding_dim) ).to(UpperCAmelCase_ )

			return {
			    "hidden_states": hidden_states,
			    "timestep": 2,
			    "proj_embedding": proj_embedding,
			    "encoder_hidden_states": encoder_hidden_states,
			}




	@property
	def       _UpperCamelCase   (			self )		->							Tuple:
			return (4, 8)




	@property
	def       _UpperCamelCase   (			self )		->							List[str]:
			return (4, 8)




	def       _UpperCamelCase   (			self )		->							List[Any]:
			lowerCamelCase					:      Optional[Any]          =   {
			    'num_attention_heads': 2,
			    'attention_head_dim': 4,
			    'num_layers': 2,
			    'embedding_dim': 8,
			    'num_embeddings': 7,
			    'additional_embeddings': 4,
			}
			lowerCamelCase					:      List[str]          =   self.dummy_input
			return init_dict, inputs_dict




	def       _UpperCamelCase   (			self )		->							int:
			lowerCamelCase					:      Union[str, Any]          =   PriorTransformer.from_pretrained(
			    'hf-internal-testing/prior-dummy'		,			output_loading_info=UpperCAmelCase_ )
			self.assertIsNotNone(UpperCAmelCase_ )
			self.assertEqual(len(loading_info['missing_keys'] )		,			0 )

			model.to(UpperCAmelCase_ )
			lowerCamelCase					:      Dict          =   model(**self.dummy_input )[0]

			assert hidden_states is not None, "Make sure output is not None"




	def       _UpperCamelCase   (			self )		->							Any:
			lowerCamelCase					:      Any          =   self.prepare_init_args_and_inputs_for_common()

			lowerCamelCase					:      Dict          =   self.model_class(**UpperCAmelCase_ )
			lowerCamelCase					:      Optional[int]          =   inspect.signature(model.forward )
			# signature.parameters is an OrderedDict => so arg_names order is deterministic
			lowerCamelCase					:      Dict          =   [*signature.parameters.keys()]

			lowerCamelCase					:      Optional[int]          =   ['hidden_states', 'timestep']
			self.assertListEqual(arg_names[:2]		,			UpperCAmelCase_ )




	def       _UpperCamelCase   (			self )		->							Union[str, Any]:
			lowerCamelCase					:      List[str]          =   PriorTransformer.from_pretrained('hf-internal-testing/prior-dummy' )
			lowerCamelCase					:      Optional[Any]          =   model.to(UpperCAmelCase_ )

			if hasattr(UpperCAmelCase_		,			'set_default_attn_processor' ):
					model.set_default_attn_processor()

			lowerCamelCase					:      str          =   self.get_dummy_seed_input()

			with torch.no_grad():
					lowerCamelCase					:      List[Any]          =   model(**UpperCAmelCase_ )[0]

			lowerCamelCase					:      List[Any]          =   output[0, :5].flatten().cpu()
			print(UpperCAmelCase_ )

			# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
			# the expected output slices are not the same for CPU and GPU.
			lowerCamelCase					:      List[str]          =   torch.tensor([-1.3436, -0.2870, 0.7538, 0.4368, -0.0239] )
			self.assertTrue(torch_all_close(UpperCAmelCase_		,			UpperCAmelCase_		,			rtol=1E-2 ) )






@slow
class 	_lowercase       (    unittest.TestCase		):




	def       _UpperCamelCase   (			self		,			UpperCAmelCase_=1		,			UpperCAmelCase_=768		,			UpperCAmelCase_=77		,			UpperCAmelCase_=0 )		->							Optional[int]:
			torch.manual_seed(UpperCAmelCase_ )
			lowerCamelCase					:      Tuple          =   batch_size
			lowerCamelCase					:      Any          =   embedding_dim
			lowerCamelCase					:      List[str]          =   num_embeddings

			lowerCamelCase					:      Optional[int]          =   torch.randn((batch_size, embedding_dim) ).to(UpperCAmelCase_ )

			lowerCamelCase					:      Union[str, Any]          =   torch.randn((batch_size, embedding_dim) ).to(UpperCAmelCase_ )
			lowerCamelCase					:      Dict          =   torch.randn((batch_size, num_embeddings, embedding_dim) ).to(UpperCAmelCase_ )

			return {
			    "hidden_states": hidden_states,
			    "timestep": 2,
			    "proj_embedding": proj_embedding,
			    "encoder_hidden_states": encoder_hidden_states,
			}




	def       _UpperCamelCase   (			self )		->							Union[str, Any]:
			# clean up the VRAM after each test
			super().tearDown()
			gc.collect()
			torch.cuda.empty_cache()




	@parameterized.expand(
	    [
	        # fmt: off
	        [13, [-0.5861,  0.1283, -0.0931,  0.0882,  0.4476,  0.1329, -0.0498,  0.0640]],
	        [37, [-0.4913,  0.0110, -0.0483,  0.0541,  0.4954, -0.0170,  0.0354,  0.1651]],
	        # fmt: on
	    ] )
	def       _UpperCamelCase   (			self		,			UpperCAmelCase_		,			UpperCAmelCase_ )		->							Dict:
			lowerCamelCase					:      Any          =   PriorTransformer.from_pretrained('kandinsky-community/kandinsky-2-1-prior'		,			subfolder='prior' )
			model.to(UpperCAmelCase_ )
			lowerCamelCase					:      Tuple          =   self.get_dummy_seed_input(seed=UpperCAmelCase_ )

			with torch.no_grad():
					lowerCamelCase					:      List[Any]          =   model(**UpperCAmelCase_ )[0]

			assert list(sample.shape ) == [1, 768]

			lowerCamelCase					:      Tuple          =   sample[0, :8].flatten().cpu()
			print(UpperCAmelCase_ )
			lowerCamelCase					:      int          =   torch.tensor(UpperCAmelCase_ )

			assert torch_all_close(UpperCAmelCase_		,			UpperCAmelCase_		,			atol=1E-3 )








"""simple docstring"""

import os
import sys
from contextlib import contextmanager


# Windows only
if os.name == "nt":
					import ctypes
					import msvcrt  # noqa






					class 	_lowercase       (    ctypes.Structure		):
						# _fields is a specific attr expected by ctypes
						lowercase_				=     [('size', ctypes.c_int), ('visible', ctypes.c_byte)]
def 		UpperCAmelCase (		):





		'''simple docstring'''


		if os.name == "nt":
				lowerCamelCase					:      Optional[int]          =   CursorInfo()
				lowerCamelCase					:      Union[str, Any]          =   ctypes.windll.kernelaa.GetStdHandle(-11      )
				ctypes.windll.kernelaa.GetConsoleCursorInfo(a_,						ctypes.byref(a_      )      )
				lowerCamelCase					:      Dict          =   False
				ctypes.windll.kernelaa.SetConsoleCursorInfo(a_,						ctypes.byref(a_      )      )
		elif os.name == "posix":
				sys.stdout.write('\033[?25l'      )
				sys.stdout.flush()
def 		UpperCAmelCase (		):





		'''simple docstring'''


		if os.name == "nt":
				lowerCamelCase					:      List[str]          =   CursorInfo()
				lowerCamelCase					:      List[Any]          =   ctypes.windll.kernelaa.GetStdHandle(-11      )
				ctypes.windll.kernelaa.GetConsoleCursorInfo(a_,						ctypes.byref(a_      )      )
				lowerCamelCase					:      Optional[Any]          =   True
				ctypes.windll.kernelaa.SetConsoleCursorInfo(a_,						ctypes.byref(a_      )      )
		elif os.name == "posix":
				sys.stdout.write('\033[?25h'      )
				sys.stdout.flush()





@contextmanager
def 		UpperCAmelCase (		):





		'''simple docstring'''


		try:
				hide_cursor()
				yield
		finally:
				show_cursor()






