




'''simple docstring'''


import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch
  from torch import nn

  from transformers import (
      MODEL_MAPPING,
      BeitForImageClassification,
      BeitForMaskedImageModeling,
      BeitForSemanticSegmentation,
      BeitModel,
  )
  from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  import PIL
  from PIL import Image

  from transformers import BeitImageProcessor





class  __lowercase    :







   def __init__(						self   ,				UpperCamelCase   ,				UpperCamelCase=100   ,				UpperCamelCase=13   ,				UpperCamelCase=30   ,				UpperCamelCase=2   ,				UpperCamelCase=3   ,				UpperCamelCase=True   ,				UpperCamelCase=True   ,				UpperCamelCase=32   ,				UpperCamelCase=4   ,				UpperCamelCase=4   ,				UpperCamelCase=37   ,				UpperCamelCase="gelu"   ,				UpperCamelCase=0.1   ,				UpperCamelCase=0.1   ,				UpperCamelCase=10   ,				UpperCamelCase=0.02   ,				UpperCamelCase=3   ,				UpperCamelCase=None   ,				UpperCamelCase=[0, 1, 2, 3]   ,				)					->     str:
     __a								=       parent
     __a								=       100
     __a								=       batch_size
     __a								=       image_size
     __a								=       patch_size
     __a								=       num_channels
     __a								=       is_training
     __a								=       use_labels
     __a								=       hidden_size
     __a								=       num_hidden_layers
     __a								=       num_attention_heads
     __a								=       intermediate_size
     __a								=       hidden_act
     __a								=       hidden_dropout_prob
     __a								=       attention_probs_dropout_prob
     __a								=       type_sequence_label_size
     __a								=       initializer_range
     __a								=       scope
     __a								=       out_indices
     __a								=       num_labels

     # in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
     __a								=       (image_size // patch_size) ** 2
     __a								=       num_patches + 1







   def 		UpperCamelCase__      (						self   )					->     List[Any]:
     __a								=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

     __a								=       None
     __a								=       None
     if self.use_labels:
       __a								=       ids_tensor([self.batch_size]   ,				self.type_sequence_label_size   )
       __a								=       ids_tensor([self.batch_size, self.image_size, self.image_size]   ,				self.num_labels   )

     __a								=       self.get_config()

     return config, pixel_values, labels, pixel_labels







   def 		UpperCamelCase__      (						self   )					->     Optional[Any]:
     return BeitConfig(
         vocab_size=self.vocab_size   ,				image_size=self.image_size   ,				patch_size=self.patch_size   ,				num_channels=self.num_channels   ,				hidden_size=self.hidden_size   ,				num_hidden_layers=self.num_hidden_layers   ,				num_attention_heads=self.num_attention_heads   ,				intermediate_size=self.intermediate_size   ,				hidden_act=self.hidden_act   ,				hidden_dropout_prob=self.hidden_dropout_prob   ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,				is_decoder=UpperCamelCase   ,				initializer_range=self.initializer_range   ,				out_indices=self.out_indices   ,				)







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     List[str]:
     __a								=       BeitModel(config=UpperCamelCase   )
     model.to(UpperCamelCase   )
     model.eval()
     __a								=       model(UpperCamelCase   )
     self.parent.assertEqual(result.last_hidden_state.shape   ,				(self.batch_size, self.seq_length, self.hidden_size)   )







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     int:
     __a								=       BeitForMaskedImageModeling(config=UpperCamelCase   )
     model.to(UpperCamelCase   )
     model.eval()
     __a								=       model(UpperCamelCase   )
     self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length - 1, self.vocab_size)   )







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     Any:
     __a								=       self.type_sequence_label_size
     __a								=       BeitForImageClassification(UpperCamelCase   )
     model.to(UpperCamelCase   )
     model.eval()
     __a								=       model(UpperCamelCase   ,				labels=UpperCamelCase   )
     self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.type_sequence_label_size)   )

     # test greyscale images
     __a								=       1
     __a								=       BeitForImageClassification(UpperCamelCase   )
     model.to(UpperCamelCase   )
     model.eval()

     __a								=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
     __a								=       model(UpperCamelCase   ,				labels=UpperCamelCase   )
     self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.type_sequence_label_size)   )







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   )					->     int:
     __a								=       self.num_labels
     __a								=       BeitForSemanticSegmentation(UpperCamelCase   )
     model.to(UpperCamelCase   )
     model.eval()
     __a								=       model(UpperCamelCase   )
     self.parent.assertEqual(
         result.logits.shape   ,				(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)   )
     __a								=       model(UpperCamelCase   ,				labels=UpperCamelCase   )
     self.parent.assertEqual(
         result.logits.shape   ,				(self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)   )







   def 		UpperCamelCase__      (						self   )					->     List[str]:
     __a								=       self.prepare_config_and_inputs()
     __a						,							__a						,							__a						,							__a								=       config_and_inputs
     __a								=       {'pixel_values': pixel_values}
     return config, inputs_dict





@require_torch
class  __lowercase    (	__magic_name__							,				__magic_name__							,				unittest.TestCase			):
   _a    						=   (
       (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
       if is_torch_available()
       else ()
   )
   _a    						=   (
       {
           """feature-extraction""": BeitModel,
           """image-classification""": BeitForImageClassification,
           """image-segmentation""": BeitForSemanticSegmentation,
       }
       if is_torch_available()
       else {}
   )

   _a    						=   False
   _a    						=   False
   _a    						=   False







   def 		UpperCamelCase__      (						self   )					->     List[Any]:
     __a								=       BeitModelTester(self   )
     __a								=       ConfigTester(self   ,				config_class=UpperCamelCase   ,				has_text_modality=UpperCamelCase   ,				hidden_size=37   )







   def 		UpperCamelCase__      (						self   )					->     List[Any]:
     self.config_tester.run_common_tests()







   @unittest.skip(reason='BEiT does not use inputs_embeds'   )
   def 		UpperCamelCase__      (						self   )					->     Tuple:
     pass







   @require_torch_multi_gpu
   @unittest.skip(reason='BEiT has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'   )
   def 		UpperCamelCase__      (						self   )					->     Optional[Any]:
     pass







   def 		UpperCamelCase__      (						self   )					->     Optional[int]:
     __a						,							__a								=       self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       __a								=       model_class(UpperCamelCase   )
       self.assertIsInstance(model.get_input_embeddings()   ,				(nn.Module)   )
       __a								=       model.get_output_embeddings()
       self.assertTrue(x is None or isinstance(UpperCamelCase   ,				nn.Linear   )   )







   def 		UpperCamelCase__      (						self   )					->     List[Any]:
     __a						,							__a								=       self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       __a								=       model_class(UpperCamelCase   )
       __a								=       inspect.signature(model.forward   )
       # signature.parameters is an OrderedDict => so arg_names order is deterministic
       __a								=       [*signature.parameters.keys()]

       __a								=       ['pixel_values']
       self.assertListEqual(arg_names[:1]   ,				UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     int:
     __a								=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     int:
     __a								=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     int:
     __a								=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     List[str]:
     __a								=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_semantic_segmentation(*UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     Optional[int]:
     if not self.model_tester.is_training:
       return

     __a						,							__a								=       self.model_tester.prepare_config_and_inputs_for_common()
     __a								=       True

     for model_class in self.all_model_classes:
       # we don't test BeitForMaskedImageModeling
       if model_class in [*get_values(UpperCamelCase   ), BeitForMaskedImageModeling]:
         continue

       __a								=       model_class(UpperCamelCase   )
       model.to(UpperCamelCase   )
       model.train()
       __a								=       self._prepare_for_class(UpperCamelCase   ,				UpperCamelCase   ,				return_labels=UpperCamelCase   )
       __a								=       model(**UpperCamelCase   ).loss
       loss.backward()







   def 		UpperCamelCase__      (						self   )					->     Dict:
     __a						,							__a								=       self.model_tester.prepare_config_and_inputs_for_common()
     if not self.model_tester.is_training:
       return

     __a								=       False
     __a								=       True

     for model_class in self.all_model_classes:
       # we don't test BeitForMaskedImageModeling
       if (
           model_class in [*get_values(UpperCamelCase   ), BeitForMaskedImageModeling]
           or not model_class.supports_gradient_checkpointing
       ):
         continue

       __a								=       model_class(UpperCamelCase   )
       model.gradient_checkpointing_enable()
       model.to(UpperCamelCase   )
       model.train()
       __a								=       self._prepare_for_class(UpperCamelCase   ,				UpperCamelCase   ,				return_labels=UpperCamelCase   )
       __a								=       model(**UpperCamelCase   ).loss
       loss.backward()







   def 		UpperCamelCase__      (						self   )					->     Union[str, Any]:
     __a						,							__a								=       self.model_tester.prepare_config_and_inputs_for_common()

     __a								=       _config_zero_init(UpperCamelCase   )
     for model_class in self.all_model_classes:
       __a								=       model_class(config=UpperCamelCase   )
       for name, param in model.named_parameters():
         # we skip lambda parameters as these require special initial values
         # determined by config.layer_scale_init_value
         if "lambda" in name:
           continue
         if param.requires_grad:
           self.assertIn(
               ((param.data.mean() * 1e9).round() / 1e9).item()   ,				[0.0, 1.0]   ,				msg=f"Parameter {name} of model {model_class} seems not properly initialized"   ,				)







   @slow
   def 		UpperCamelCase__      (						self   )					->     str:
     for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       __a								=       BeitModel.from_pretrained(UpperCamelCase   )
       self.assertIsNotNone(UpperCamelCase   )


def     SCREAMING_SNAKE_CASE							(   ):
  __a								=       Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'      )
  return image





@require_torch
@require_vision
class  __lowercase    (	unittest.TestCase			):







   @cached_property
   def 		UpperCamelCase__      (						self   )					->     List[str]:
     return BeitImageProcessor.from_pretrained('microsoft/beit-base-patch16-224'   ) if is_vision_available() else None







   @slow
   def 		UpperCamelCase__      (						self   )					->     int:
     __a								=       BeitForMaskedImageModeling.from_pretrained('microsoft/beit-base-patch16-224-pt22k'   ).to(UpperCamelCase   )

     __a								=       self.default_image_processor
     __a								=       prepare_img()
     __a								=       image_processor(images=UpperCamelCase   ,				return_tensors='pt'   ).pixel_values.to(UpperCamelCase   )

     # prepare bool_masked_pos
     __a								=       torch.ones((1, 196)   ,				dtype=torch.bool   ).to(UpperCamelCase   )

     # forward pass
     with torch.no_grad():
       __a								=       model(pixel_values=UpperCamelCase   ,				bool_masked_pos=UpperCamelCase   )
     __a								=       outputs.logits

     # verify the logits
     __a								=       torch.Size((1, 196, 8192)   )
     self.assertEqual(logits.shape   ,				UpperCamelCase   )

     __a								=       torch.tensor(
         [[-3.2_437, 0.5_072, -13.9_174], [-3.2_456, 0.4_948, -13.9_401], [-3.2_033, 0.5_121, -13.8_550]]   ).to(UpperCamelCase   )

     self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]   ,				UpperCamelCase   ,				atol=1e-2   )   )







   @slow
   def 		UpperCamelCase__      (						self   )					->     Optional[int]:
     __a								=       BeitForImageClassification.from_pretrained('microsoft/beit-base-patch16-224'   ).to(UpperCamelCase   )

     __a								=       self.default_image_processor
     __a								=       prepare_img()
     __a								=       image_processor(images=UpperCamelCase   ,				return_tensors='pt'   ).to(UpperCamelCase   )

     # forward pass
     with torch.no_grad():
       __a								=       model(**UpperCamelCase   )
     __a								=       outputs.logits

     # verify the logits
     __a								=       torch.Size((1, 1000)   )
     self.assertEqual(logits.shape   ,				UpperCamelCase   )

     __a								=       torch.tensor([-1.2_385, -1.0_987, -1.0_108]   ).to(UpperCamelCase   )

     self.assertTrue(torch.allclose(logits[0, :3]   ,				UpperCamelCase   ,				atol=1e-4   )   )

     __a								=       281
     self.assertEqual(logits.argmax(-1   ).item()   ,				UpperCamelCase   )







   @slow
   def 		UpperCamelCase__      (						self   )					->     Any:
     __a								=       BeitForImageClassification.from_pretrained('microsoft/beit-large-patch16-224-pt22k-ft22k'   ).to(
         UpperCamelCase   )

     __a								=       self.default_image_processor
     __a								=       prepare_img()
     __a								=       image_processor(images=UpperCamelCase   ,				return_tensors='pt'   ).to(UpperCamelCase   )

     # forward pass
     with torch.no_grad():
       __a								=       model(**UpperCamelCase   )
     __a								=       outputs.logits

     # verify the logits
     __a								=       torch.Size((1, 2_1841)   )
     self.assertEqual(logits.shape   ,				UpperCamelCase   )

     __a								=       torch.tensor([1.6_881, -0.2_787, 0.5_901]   ).to(UpperCamelCase   )

     self.assertTrue(torch.allclose(logits[0, :3]   ,				UpperCamelCase   ,				atol=1e-4   )   )

     __a								=       2396
     self.assertEqual(logits.argmax(-1   ).item()   ,				UpperCamelCase   )







   @slow
   def 		UpperCamelCase__      (						self   )					->     int:
     __a								=       BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640'   )
     __a								=       model.to(UpperCamelCase   )

     __a								=       BeitImageProcessor(do_resize=UpperCamelCase   ,				size=640   ,				do_center_crop=UpperCamelCase   )

     __a								=       load_dataset('hf-internal-testing/fixtures_ade20k'   ,				split='test'   )
     __a								=       Image.open(ds[0]['file']   )
     __a								=       image_processor(images=UpperCamelCase   ,				return_tensors='pt'   ).to(UpperCamelCase   )

     # forward pass
     with torch.no_grad():
       __a								=       model(**UpperCamelCase   )
     __a								=       outputs.logits

     # verify the logits
     __a								=       torch.Size((1, 150, 160, 160)   )
     self.assertEqual(logits.shape   ,				UpperCamelCase   )

     __a								=       version.parse(PIL.__version__   ) < version.parse('9.0.0'   )

     if is_pillow_less_than_a:
       __a								=       torch.tensor(
           [
               [[-4.9_225, -2.3_954, -3.0_522], [-2.8_822, -1.0_046, -1.7_561], [-2.9_549, -1.3_228, -2.1_347]],
               [[-5.8_168, -3.4_129, -4.0_778], [-3.8_651, -2.2_214, -3.0_277], [-3.8_356, -2.4_643, -3.3_535]],
               [[-0.0_078, 3.9_952, 4.0_754], [2.9_856, 4.6_944, 5.0_035], [3.2_413, 4.7_813, 4.9_969]],
           ]   ,				device=UpperCamelCase   ,				)
     else:
       __a								=       torch.tensor(
           [
               [[-4.8_960, -2.3_688, -3.0_355], [-2.8_478, -0.9_836, -1.7_418], [-2.9_449, -1.3_332, -2.1_456]],
               [[-5.8_081, -3.4_124, -4.1_006], [-3.8_561, -2.2_081, -3.0_323], [-3.8_365, -2.4_601, -3.3_669]],
               [[-0.0_309, 3.9_868, 4.0_540], [2.9_640, 4.6_877, 4.9_976], [3.2_081, 4.7_690, 4.9_942]],
           ]   ,				device=UpperCamelCase   ,				)

     self.assertTrue(torch.allclose(logits[0, :3, :3, :3]   ,				UpperCamelCase   ,				atol=1e-4   )   )







   @slow
   def 		UpperCamelCase__      (						self   )					->     Any:
     __a								=       BeitForSemanticSegmentation.from_pretrained('microsoft/beit-base-finetuned-ade-640-640'   )
     __a								=       model.to(UpperCamelCase   )

     __a								=       BeitImageProcessor(do_resize=UpperCamelCase   ,				size=640   ,				do_center_crop=UpperCamelCase   )

     __a								=       load_dataset('hf-internal-testing/fixtures_ade20k'   ,				split='test'   )
     __a								=       Image.open(ds[0]['file']   )
     __a								=       image_processor(images=UpperCamelCase   ,				return_tensors='pt'   ).to(UpperCamelCase   )

     # forward pass
     with torch.no_grad():
       __a								=       model(**UpperCamelCase   )

     __a								=       outputs.logits.detach().cpu()

     __a								=       image_processor.post_process_semantic_segmentation(outputs=UpperCamelCase   ,				target_sizes=[(500, 300)]   )
     __a								=       torch.Size((500, 300)   )
     self.assertEqual(segmentation[0].shape   ,				UpperCamelCase   )

     __a								=       image_processor.post_process_semantic_segmentation(outputs=UpperCamelCase   )
     __a								=       torch.Size((160, 160)   )
     self.assertEqual(segmentation[0].shape   ,				UpperCamelCase   )









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







   @property
   def 		UpperCamelCase__      (						self   )					->     int:
     __a								=       4
     __a								=       8
     __a								=       7

     __a								=       floats_tensor((batch_size, embedding_dim)   ).to(UpperCamelCase   )

     __a								=       floats_tensor((batch_size, embedding_dim)   ).to(UpperCamelCase   )
     __a								=       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   )					->     Union[str, Any]:
     torch.manual_seed(UpperCamelCase   )
     __a								=       4
     __a								=       8
     __a								=       7

     __a								=       torch.randn((batch_size, embedding_dim)   ).to(UpperCamelCase   )

     __a								=       torch.randn((batch_size, embedding_dim)   ).to(UpperCamelCase   )
     __a								=       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   )					->     List[str]:
     return (4, 8)







   @property
   def 		UpperCamelCase__      (						self   )					->     str:
     return (4, 8)







   def 		UpperCamelCase__      (						self   )					->     Dict:
     __a								=       {
         'num_attention_heads': 2,
         'attention_head_dim': 4,
         'num_layers': 2,
         'embedding_dim': 8,
         'num_embeddings': 7,
         'additional_embeddings': 4,
     }
     __a								=       self.dummy_input
     return init_dict, inputs_dict







   def 		UpperCamelCase__      (						self   )					->     List[Any]:
     __a						,							__a								=       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   )
     __a								=       model(**self.dummy_input   )[0]

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







   def 		UpperCamelCase__      (						self   )					->     Optional[Any]:
     __a						,							__a								=       self.prepare_init_args_and_inputs_for_common()

     __a								=       self.model_class(**UpperCamelCase   )
     __a								=       inspect.signature(model.forward   )
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     __a								=       [*signature.parameters.keys()]

     __a								=       ['hidden_states', 'timestep']
     self.assertListEqual(arg_names[:2]   ,				UpperCamelCase   )







   def 		UpperCamelCase__      (						self   )					->     str:
     __a								=       PriorTransformer.from_pretrained('hf-internal-testing/prior-dummy'   )
     __a								=       model.to(UpperCamelCase   )

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

     __a								=       self.get_dummy_seed_input()

     with torch.no_grad():
       __a								=       model(**UpperCamelCase   )[0]

     __a								=       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.
     __a								=       torch.tensor([-1.3_436, -0.2_870, 0.7_538, 0.4_368, -0.0_239]   )
     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   )					->     List[str]:
     torch.manual_seed(UpperCamelCase   )
     __a								=       batch_size
     __a								=       embedding_dim
     __a								=       num_embeddings

     __a								=       torch.randn((batch_size, embedding_dim)   ).to(UpperCamelCase   )

     __a								=       torch.randn((batch_size, embedding_dim)   ).to(UpperCamelCase   )
     __a								=       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   )					->     List[str]:
     # 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 		UpperCamelCase__      (						self   ,				UpperCamelCase   ,				UpperCamelCase   )					->     Optional[Any]:
     __a								=       PriorTransformer.from_pretrained('kandinsky-community/kandinsky-2-1-prior'   ,				subfolder='prior'   )
     model.to(UpperCamelCase   )
     __a								=       self.get_dummy_seed_input(seed=UpperCamelCase   )

     with torch.no_grad():
       __a								=       model(**UpperCamelCase   )[0]

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

     __a								=       sample[0, :8].flatten().cpu()
     print(UpperCamelCase   )
     __a								=       torch.tensor(UpperCamelCase   )

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



