






'''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  snake_case__      (					snake_case__     ,				unittest.TestCase):
					a_        =					PriorTransformer
					a_        =					"hidden_states"




					@property
					def 		A      (  self  :		Any      )    ->							List[Any]:
								UpperCAmelCase_      :       int          =       4
								UpperCAmelCase_      :       Dict          =       8
								UpperCAmelCase_      :       List[str]          =       7

								UpperCAmelCase_      :       Any          =       floats_tensor((batch_size, embedding_dim)      ).to(__snake_case      )

								UpperCAmelCase_      :       Union[str, Any]          =       floats_tensor((batch_size, embedding_dim)      ).to(__snake_case      )
								UpperCAmelCase_      :       Dict          =       floats_tensor((batch_size, num_embeddings, embedding_dim)      ).to(__snake_case      )

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




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

								UpperCAmelCase_      :       Tuple          =       torch.randn((batch_size, embedding_dim)      ).to(__snake_case      )

								UpperCAmelCase_      :       List[str]          =       torch.randn((batch_size, embedding_dim)      ).to(__snake_case      )
								UpperCAmelCase_      :       Optional[int]          =       torch.randn((batch_size, num_embeddings, embedding_dim)      ).to(__snake_case      )

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




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




					@property
					def 		A      (  self  :		List[Any]      )    ->							Dict:
								return (4, 8)




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




					def 		A      (  self  :		int      )    ->							Tuple:
								UpperCAmelCase_    ,      UpperCAmelCase_      :       Optional[Any]          =       PriorTransformer.from_pretrained(
								    '''hf-internal-testing/prior-dummy'''				, output_loading_info=__snake_case      )
								self.assertIsNotNone(__snake_case      )
								self.assertEqual(len(loading_info['''missing_keys''']      )				, 0      )

								model.to(__snake_case      )
								UpperCAmelCase_      :       Optional[Any]          =       model(**self.dummy_input      )[0]

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




					def 		A      (  self  :		List[str]      )    ->							Any:
								UpperCAmelCase_    ,      UpperCAmelCase_      :       int          =       self.prepare_init_args_and_inputs_for_common()

								UpperCAmelCase_      :       Tuple          =       self.model_class(**__snake_case      )
								UpperCAmelCase_      :       Union[str, Any]          =       inspect.signature(model.forward      )
								# signature.parameters is an OrderedDict => so arg_names order is deterministic
								UpperCAmelCase_      :       List[str]          =       [*signature.parameters.keys()]

								UpperCAmelCase_      :       Dict          =       ['''hidden_states''', '''timestep''']
								self.assertListEqual(arg_names[:2]				, __snake_case      )





					def 		A      (  self  :		Dict      )    ->							List[str]:
								UpperCAmelCase_      :       Optional[Any]          =       PriorTransformer.from_pretrained('''hf-internal-testing/prior-dummy'''      )
								UpperCAmelCase_      :       Tuple          =       model.to(__snake_case      )

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

								UpperCAmelCase_      :       str          =       self.get_dummy_seed_input()

								with torch.no_grad():
											UpperCAmelCase_      :       Optional[Any]          =       model(**__snake_case      )[0]

								UpperCAmelCase_      :       Optional[int]          =       output[0, :5].flatten().cpu()
								print(__snake_case      )

								# 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.
								UpperCAmelCase_      :       Optional[int]          =       torch.tensor([-1.3_436, -0.2_870, 0.7_538, 0.4_368, -0.0_239]      )
								self.assertTrue(torch_all_close(__snake_case				, __snake_case				, rtol=1e-2      )      )






@slow
class  snake_case__      (					unittest.TestCase):




					def 		A      (  self  :		Tuple				, _A  :		Optional[int]=1				, _A  :		Dict=7_68				, _A  :		Union[str, Any]=77				, _A  :		int=0      )    ->							int:
								torch.manual_seed(__snake_case      )
								UpperCAmelCase_      :       List[Any]          =       batch_size
								UpperCAmelCase_      :       str          =       embedding_dim
								UpperCAmelCase_      :       List[Any]          =       num_embeddings

								UpperCAmelCase_      :       str          =       torch.randn((batch_size, embedding_dim)      ).to(__snake_case      )

								UpperCAmelCase_      :       int          =       torch.randn((batch_size, embedding_dim)      ).to(__snake_case      )
								UpperCAmelCase_      :       Any          =       torch.randn((batch_size, num_embeddings, embedding_dim)      ).to(__snake_case      )

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




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





					@parameterized.expand(
					    [
					        # fmt: off
					        [13, [-0.5_861,  0.1_283, -0.0_931,  0.0_882,  0.4_476,  0.1_329, -0.0_498,  0.0_640]],
					        [37, [-0.4_913,  0.0_110, -0.0_483,  0.0_541,  0.4_954, -0.0_170,  0.0_354,  0.1_651]],
					        # fmt: on
					    ]      )
					def 		A      (  self  :		Optional[Any]				, _A  :		Union[str, Any]				, _A  :		str      )    ->							str:
								UpperCAmelCase_      :       Dict          =       PriorTransformer.from_pretrained('''kandinsky-community/kandinsky-2-1-prior'''				, subfolder='''prior'''      )
								model.to(__snake_case      )
								UpperCAmelCase_      :       Tuple          =       self.get_dummy_seed_input(seed=__snake_case      )

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

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

								UpperCAmelCase_      :       List[str]          =       sample[0, :8].flatten().cpu()
								print(__snake_case      )
								UpperCAmelCase_      :       Optional[int]          =       torch.tensor(__snake_case      )

								assert torch_all_close(__snake_case				, __snake_case				, atol=1e-3      )






from functools import lru_cache
def      UpperCAmelCase__		(UpperCamelCase_							):



     """simple docstring"""






     snake_case											=							2
     snake_case											=							set()
     while i * i <= n:
          if n % i:
               i += 1
          else:
               n //= i
               factors.add(UpperCamelCase_							)
     if n > 1:
          factors.add(UpperCamelCase_							)
     return factors
@lru_cache
def      UpperCAmelCase__		(UpperCamelCase_							):



     """simple docstring"""






     return len(unique_prime_factors(UpperCamelCase_							)							)
def      UpperCAmelCase__		(UpperCamelCase_							):



     """simple docstring"""






     return len(set(UpperCamelCase_							)							) in (0, 1)
def      UpperCAmelCase__		(UpperCamelCase_							):



     """simple docstring"""






     snake_case											=							2

     while True:
          # Increment each value of a generated range
          snake_case											=							[base + i for i in range(UpperCamelCase_							)]

          # Run elements through out unique_prime_factors function
          # Append our target number to the end.
          snake_case											=							[upf_len(UpperCamelCase_							) for x in group]
          checker.append(UpperCamelCase_							)

          # If all numbers in the list are equal, return the group variable.
          if equality(UpperCamelCase_							):
               return group

          # Increment our base variable by 1
          base += 1
def      UpperCAmelCase__		(UpperCamelCase_ = 4							):



     """simple docstring"""






     snake_case											=							run(UpperCamelCase_							)
     return results[0] if len(UpperCamelCase_							) else None


if __name__ == "__main__":
 print(solution())




