

"""simple docstring"""
import argparse

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    CLIPTokenizer,
    CLIPTokenizerFast,
    VideoMAEImageProcessor,
    XCLIPConfig,
    XCLIPModel,
    XCLIPProcessor,
    XCLIPTextConfig,
    XCLIPVisionConfig,
)







def 						__a       (     _SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE       )     ->str:
  a__:					Tuple      							=     XCLIPTextConfig()

  # derive patch size from model name
  a__:					Dict      							=     model_name.find('patch'       )
  a__:					Dict      							=     int(model_name[start_idx + len('patch'       ) : start_idx + len('patch'       ) + 2]       )
  a__:					Optional[Any]      							=     XCLIPVisionConfig(patch_size=_SCREAMING_SNAKE_CASE		,					num_frames=_SCREAMING_SNAKE_CASE       )

  if "large" in model_name:
    a__:					Any      							=     768
    a__:					Any      							=     3072
    a__:					List[Any]      							=     12

    a__:					Any      							=     1024
    a__:					List[Any]      							=     4096
    a__:					Union[str, Any]      							=     16
    a__:					int      							=     24
    a__:					Any      							=     768
    a__:					List[str]      							=     3072

  if model_name == "xclip-large-patch14-16-frames":
    a__:					str      							=     336

  a__:					int      							=     XCLIPConfig.from_text_vision_configs(_SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE       )

  if "large" in model_name:
    a__:					Union[str, Any]      							=     768

  return config







def 						__a       (     _SCREAMING_SNAKE_CASE       )     ->Optional[int]:
  # text encoder
  if name == "token_embedding.weight":
    a__:					Dict      							=     name.replace('token_embedding.weight'		,					'text_model.embeddings.token_embedding.weight'       )
  if name == "positional_embedding":
    a__:					Dict      							=     name.replace('positional_embedding'		,					'text_model.embeddings.position_embedding.weight'       )
  if "ln_1" in name:
    a__:					List[Any]      							=     name.replace('ln_1'		,					'layer_norm1'       )
  if "ln_2" in name:
    a__:					int      							=     name.replace('ln_2'		,					'layer_norm2'       )
  if "c_fc" in name:
    a__:					Optional[Any]      							=     name.replace('c_fc'		,					'fc1'       )
  if "c_proj" in name:
    a__:					Any      							=     name.replace('c_proj'		,					'fc2'       )
  if name.startswith('transformer.resblocks'       ):
    a__:					Any      							=     name.replace('transformer.resblocks'		,					'text_model.encoder.layers'       )
  if "attn.out_proj" in name and "message" not in name:
    a__:					List[Any]      							=     name.replace('attn.out_proj'		,					'self_attn.out_proj'       )
  if "ln_final" in name:
    a__:					Optional[Any]      							=     name.replace('ln_final'		,					'text_model.final_layer_norm'       )
  # visual encoder
  if name == "visual.class_embedding":
    a__:					Dict      							=     name.replace('visual.class_embedding'		,					'vision_model.embeddings.class_embedding'       )
  if name == "visual.positional_embedding":
    a__:					Tuple      							=     name.replace('visual.positional_embedding'		,					'vision_model.embeddings.position_embedding.weight'       )
  if name.startswith('visual.transformer.resblocks'       ):
    a__:					Tuple      							=     name.replace('visual.transformer.resblocks'		,					'vision_model.encoder.layers'       )
  if "visual.conv1" in name:
    a__:					Optional[Any]      							=     name.replace('visual.conv1'		,					'vision_model.embeddings.patch_embedding'       )
  if "visual.ln_pre" in name:
    a__:					Any      							=     name.replace('visual.ln_pre'		,					'vision_model.pre_layernorm'       )
  if "visual.ln_post" in name:
    a__:					Union[str, Any]      							=     name.replace('visual.ln_post'		,					'vision_model.post_layernorm'       )
  if "visual.proj" in name:
    a__:					Optional[Any]      							=     name.replace('visual.proj'		,					'visual_projection.weight'       )
  if "text_projection" in name:
    a__:					List[str]      							=     name.replace('text_projection'		,					'text_projection.weight'       )
  # things on top
  if "prompts_visual_proj" in name:
    a__:					Union[str, Any]      							=     name.replace('prompts_visual_proj'		,					'prompts_visual_projection'       )
  if "prompts_visual_ln" in name:
    a__:					Optional[Any]      							=     name.replace('prompts_visual_ln'		,					'prompts_visual_layernorm'       )
  # mit
  if name == "mit.positional_embedding":
    a__:					Dict      							=     name.replace('positional'		,					'position'       )
  if name.startswith('mit.resblocks'       ):
    a__:					Tuple      							=     name.replace('mit.resblocks'		,					'mit.encoder.layers'       )
  # prompts generator
  if name.startswith('prompts_generator.norm'       ):
    a__:					Optional[Any]      							=     name.replace('prompts_generator.norm'		,					'prompts_generator.layernorm'       )

  return name







def 						__a       (     _SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE       )     ->str:
  for key in orig_state_dict.copy().keys():
    a__:					int      							=     orig_state_dict.pop(_SCREAMING_SNAKE_CASE       )

    if "attn.in_proj" in key:
      a__:					Optional[Any]      							=     key.split('.'       )
      if key.startswith('visual'       ):
        a__:					Optional[Any]      							=     key_split[3]
        a__:					str      							=     config.vision_config.hidden_size
        if "message_attn" in key:
          if "weight" in key:
            a__:					Dict      							=     val[
                :dim, :
            ]
            a__:					List[str]      							=     val[
                dim : dim * 2, :
            ]
            a__:					List[Any]      							=     val[
                -dim:, :
            ]
          else:
            a__:					Union[str, Any]      							=     val[
                :dim
            ]
            a__:					Union[str, Any]      							=     val[
                dim : dim * 2
            ]
            a__:					str      							=     val[
                -dim:
            ]
        else:
          if "weight" in key:
            a__:					List[Any]      							=     val[
                :dim, :
            ]
            a__:					Dict      							=     val[
                dim : dim * 2, :
            ]
            a__:					List[Any]      							=     val[
                -dim:, :
            ]
          else:
            a__:					Any      							=     val[:dim]
            a__:					Dict      							=     val[
                dim : dim * 2
            ]
            a__:					List[str]      							=     val[-dim:]
      elif key.startswith('mit'       ):
        a__:					int      							=     key_split[2]
        a__:					Optional[Any]      							=     config.vision_config.mit_hidden_size
        if "weight" in key:
          a__:					Optional[Any]      							=     val[:dim, :]
          a__:					Optional[Any]      							=     val[dim : dim * 2, :]
          a__:					int      							=     val[-dim:, :]
        else:
          a__:					str      							=     val[:dim]
          a__:					Optional[int]      							=     val[dim : dim * 2]
          a__:					Optional[int]      							=     val[-dim:]
      else:
        a__:					int      							=     key_split[2]
        a__:					List[Any]      							=     config.text_config.hidden_size
        if "weight" in key:
          a__:					Tuple      							=     val[:dim, :]
          a__:					Optional[int]      							=     val[
              dim : dim * 2, :
          ]
          a__:					Optional[int]      							=     val[-dim:, :]
        else:
          a__:					List[Any]      							=     val[:dim]
          a__:					List[Any]      							=     val[
              dim : dim * 2
          ]
          a__:					int      							=     val[-dim:]
    else:
      a__:					Optional[int]      							=     rename_key(_SCREAMING_SNAKE_CASE       )
      if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
        a__:					Optional[Any]      							=     val.T
      a__:					List[Any]      							=     val

  return orig_state_dict







def 						__a       (     _SCREAMING_SNAKE_CASE       )     ->List[str]:
  if num_frames == 8:
    a__:					List[Any]      							=     'eating_spaghetti_8_frames.npy'
  elif num_frames == 16:
    a__:					List[Any]      							=     'eating_spaghetti.npy'
  elif num_frames == 32:
    a__:					Optional[Any]      							=     'eating_spaghetti_32_frames.npy'
  a__:					str      							=     hf_hub_download(
      repo_id='hf-internal-testing/spaghetti-video'		,					filename=_SCREAMING_SNAKE_CASE		,					repo_type='dataset'		,					)
  a__:					Optional[Any]      							=     np.load(_SCREAMING_SNAKE_CASE       )
  return list(_SCREAMING_SNAKE_CASE       )







def 						__a       (     _SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE=None		,					_SCREAMING_SNAKE_CASE=False       )     ->str:
  a__:					Union[str, Any]      							=     {
      # fully supervised kinetics-400 checkpoints
      'xclip-base-patch32': 'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_8.pth',
      'xclip-base-patch32-16-frames': (
          'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_16.pth'
      ),
      'xclip-base-patch16': 'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_8.pth',
      'xclip-base-patch16-16-frames': (
          'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_16.pth'
      ),
      'xclip-large-patch14': 'https://drive.google.com/u/0/uc?id=1NUOImq0o5DlQTST17iIP3vG7DgmHQuCx&amp;export=download&amp;confirm=t&amp;uuid=b26caedc-88e2-473e-830a-9d158b653cdb',
      'xclip-large-patch14-16-frames': 'https://drive.google.com/u/0/uc?id=1FOYgnJc097OJ4lGwtRCCydQyVPJEOH7d&amp;export=download&amp;confirm=t&amp;uuid=538fa810-e671-4050-b385-9a623f89804f',
      # fully supervised kinetics-600 checkpoints
      'xclip-base-patch16-kinetics-600': (
          'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_8.pth'
      ),
      'xclip-base-patch16-kinetics-600-16-frames': (
          'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_16.pth'
      ),
      'xclip-large-patch14-kinetics-600': 'https://drive.google.com/u/0/uc?id=1FV8C1INuM91sLAN4ImjzePLIlpMSihwV&amp;export=download&amp;confirm=t&amp;uuid=141d4977-4a65-44ae-864f-4b0c19f838be',
      # few shot
      'xclip-base-patch16-hmdb-2-shot': (
          'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_2.pth'
      ),
      'xclip-base-patch16-hmdb-4-shot': (
          'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_4.pth'
      ),
      'xclip-base-patch16-hmdb-8-shot': (
          'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_8.pth'
      ),
      'xclip-base-patch16-hmdb-16-shot': (
          'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_16.pth'
      ),
      'xclip-base-patch16-ucf-2-shot': (
          'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_2.pth'
      ),
      'xclip-base-patch16-ucf-4-shot': (
          'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_4.pth'
      ),
      'xclip-base-patch16-ucf-8-shot': (
          'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_8.pth'
      ),
      'xclip-base-patch16-ucf-16-shot': (
          'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_16.pth'
      ),
      # zero shot
      'xclip-base-patch16-zero-shot': 'https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/zero.pth',
  }

  a__:					Any      							=     model_to_url[model_name]
  a__:					Tuple      							=     8
  if "16-frames" in model_name:
    a__:					str      							=     16
  elif "shot" in model_name:
    a__:					int      							=     32

  a__:					Tuple      							=     get_xclip_config(_SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE       )
  a__:					List[str]      							=     XCLIPModel(_SCREAMING_SNAKE_CASE       )
  model.eval()

  if "drive" in checkpoint_url:
    a__:					Optional[Any]      							=     'pytorch_model.bin'
    gdown.cached_download(_SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE		,					quiet=_SCREAMING_SNAKE_CASE       )
    a__:					Optional[int]      							=     torch.load(_SCREAMING_SNAKE_CASE		,					map_location='cpu'       )['model']
  else:
    a__:					List[Any]      							=     torch.hub.load_state_dict_from_url(_SCREAMING_SNAKE_CASE       )['model']

  a__:					int      							=     convert_state_dict(_SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE       )

  a__:					int      							=     XCLIPModel(_SCREAMING_SNAKE_CASE       )
  a__:					Optional[int]      							=     model.load_state_dict(_SCREAMING_SNAKE_CASE		,					strict=_SCREAMING_SNAKE_CASE       )
  assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
  model.eval()

  a__:					Optional[Any]      							=     336 if model_name == 'xclip-large-patch14-16-frames' else 224
  a__:					List[str]      							=     VideoMAEImageProcessor(size=_SCREAMING_SNAKE_CASE       )
  a__:					Optional[int]      							=     CLIPTokenizer.from_pretrained('openai/clip-vit-base-patch32'       )
  a__:					List[str]      							=     CLIPTokenizerFast.from_pretrained('openai/clip-vit-base-patch32'       )
  a__:					Optional[int]      							=     XCLIPProcessor(image_processor=_SCREAMING_SNAKE_CASE		,					tokenizer=_SCREAMING_SNAKE_CASE       )

  a__:					str      							=     prepare_video(_SCREAMING_SNAKE_CASE       )
  a__:					Any      							=     processor(
      text=['playing sports', 'eating spaghetti', 'go shopping']		,					videos=_SCREAMING_SNAKE_CASE		,					return_tensors='pt'		,					padding=_SCREAMING_SNAKE_CASE       )

  print('Shape of pixel values:'		,					inputs.pixel_values.shape       )

  with torch.no_grad():
    a__:					Optional[int]      							=     model(**_SCREAMING_SNAKE_CASE       )

  # Verify outputs
  a__:					int      							=     outputs.logits_per_video
  a__:					Optional[int]      							=     logits_per_video.softmax(dim=1       )
  print('Probs:'		,					_SCREAMING_SNAKE_CASE       )
  # kinetics-400
  if model_name == "xclip-base-patch32":
    a__:					Union[str, Any]      							=     torch.tensor([[0.0_019, 0.9_951, 0.0_030]]       )
  elif model_name == "xclip-base-patch32-16-frames":
    a__:					List[Any]      							=     torch.tensor([[7.0999e-04, 9.9883e-01, 4.5580e-04]]       )
  elif model_name == "xclip-base-patch16":
    a__:					List[Any]      							=     torch.tensor([[0.0_083, 0.9_681, 0.0_236]]       )
  elif model_name == "xclip-base-patch16-16-frames":
    a__:					int      							=     torch.tensor([[7.6937e-04, 9.9728e-01, 1.9473e-03]]       )
  elif model_name == "xclip-large-patch14":
    a__:					List[str]      							=     torch.tensor([[0.0_062, 0.9_864, 0.0_075]]       )
  elif model_name == "xclip-large-patch14-16-frames":
    a__:					Optional[int]      							=     torch.tensor([[3.3877e-04, 9.9937e-01, 2.8888e-04]]       )
  # kinetics-600
  elif model_name == "xclip-base-patch16-kinetics-600":
    a__:					int      							=     torch.tensor([[0.0_555, 0.8_914, 0.0_531]]       )
  elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
    a__:					Any      							=     torch.tensor([[3.8554e-04, 9.9929e-01, 3.2754e-04]]       )
  elif model_name == "xclip-large-patch14-kinetics-600":
    a__:					Dict      							=     torch.tensor([[0.0_036, 0.9_920, 0.0_045]]       )
  # few shot
  elif model_name == "xclip-base-patch16-hmdb-2-shot":
    a__:					Tuple      							=     torch.tensor([[7.1890e-06, 9.9994e-01, 5.6559e-05]]       )
  elif model_name == "xclip-base-patch16-hmdb-4-shot":
    a__:					List[Any]      							=     torch.tensor([[1.0320e-05, 9.9993e-01, 6.2435e-05]]       )
  elif model_name == "xclip-base-patch16-hmdb-8-shot":
    a__:					Optional[Any]      							=     torch.tensor([[4.1377e-06, 9.9990e-01, 9.8386e-05]]       )
  elif model_name == "xclip-base-patch16-hmdb-16-shot":
    a__:					int      							=     torch.tensor([[4.1347e-05, 9.9962e-01, 3.3411e-04]]       )
  elif model_name == "xclip-base-patch16-ucf-2-shot":
    a__:					List[str]      							=     torch.tensor([[8.5857e-05, 9.9928e-01, 6.3291e-04]]       )
  elif model_name == "xclip-base-patch16-ucf-4-shot":
    a__:					List[str]      							=     torch.tensor([[8.5857e-05, 9.9928e-01, 6.3291e-04]]       )
  elif model_name == "xclip-base-patch16-ucf-8-shot":
    a__:					List[Any]      							=     torch.tensor([[0.0_027, 0.9_904, 0.0_070]]       )
  elif model_name == "xclip-base-patch16-ucf-16-shot":
    a__:					List[Any]      							=     torch.tensor([[9.8219e-04, 9.9593e-01, 3.0863e-03]]       )
  # zero shot
  elif model_name == "xclip-base-patch16-zero-shot":
    a__:					str      							=     torch.tensor([[3.5082e-04, 9.9785e-01, 1.7966e-03]]       )
  else:
    raise ValueError(F'Model name {model_name} not supported'       )
  assert torch.allclose(_SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE		,					atol=1e-3       )
  print('Looks ok!'       )

  if pytorch_dump_folder_path is not None:
    print(F'Saving model {model_name} to {pytorch_dump_folder_path}'       )
    model.save_pretrained(_SCREAMING_SNAKE_CASE       )

  if push_to_hub:
    print('Pushing model, processor and slow tokenizer files to the hub...'       )
    model.push_to_hub(_SCREAMING_SNAKE_CASE		,					organization='nielsr'       )
    processor.push_to_hub(_SCREAMING_SNAKE_CASE		,					organization='nielsr'       )
    slow_tokenizer.push_to_hub(_SCREAMING_SNAKE_CASE		,					organization='nielsr'       )


if __name__ == "__main__":
 lowercase__ 	=     argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '--model_name',
     default='xclip-base-patch32',
     type=str,
     help='Name of the model.',
 )
 parser.add_argument(
     '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
 )
 parser.add_argument(
     '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
 )

 lowercase__ 	=     parser.parse_args()
 convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)



"""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	(					__lowerCAmelCase	,				unittest.TestCase    ):
      a__				=		PriorTransformer
      a__				=		"""hidden_states"""



      @property
      def   lowerCamelCase_       (						self)     ->      Tuple:





        '''simple docstring'''





        a__:					Union[str, Any]      							=     4
        a__:					Any      							=     8
        a__:					Optional[Any]      							=     7

        a__:					Tuple      							=     floats_tensor((batch_size, embedding_dim)).to(lowercase)

        a__:					Optional[int]      							=     floats_tensor((batch_size, embedding_dim)).to(lowercase)
        a__:					List[str]      							=     floats_tensor((batch_size, num_embeddings, embedding_dim)).to(lowercase)

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



      def   lowerCamelCase_       (						self   ,      lowercase=0)     ->      str:





        '''simple docstring'''





        torch.manual_seed(lowercase)
        a__:					Optional[Any]      							=     4
        a__:					Optional[Any]      							=     8
        a__:					Union[str, Any]      							=     7

        a__:					Optional[Any]      							=     torch.randn((batch_size, embedding_dim)).to(lowercase)

        a__:					List[str]      							=     torch.randn((batch_size, embedding_dim)).to(lowercase)
        a__:					Tuple      							=     torch.randn((batch_size, num_embeddings, embedding_dim)).to(lowercase)

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



      @property
      def   lowerCamelCase_       (						self)     ->      str:





        '''simple docstring'''





        return (4, 8)



      @property
      def   lowerCamelCase_       (						self)     ->      Optional[int]:





        '''simple docstring'''





        return (4, 8)



      def   lowerCamelCase_       (						self)     ->      str:





        '''simple docstring'''





        a__:					int      							=     {
            'num_attention_heads': 2,
            'attention_head_dim': 4,
            'num_layers': 2,
            'embedding_dim': 8,
            'num_embeddings': 7,
            'additional_embeddings': 4,
        }
        a__:					Union[str, Any]      							=     self.dummy_input
        return init_dict, inputs_dict



      def   lowerCamelCase_       (						self)     ->      Union[str, Any]:





        '''simple docstring'''





        a__							,   a__:					Union[str, Any]      							=     PriorTransformer.from_pretrained(
            'hf-internal-testing/prior-dummy'   ,      output_loading_info=lowercase)
        self.assertIsNotNone(lowercase)
        self.assertEqual(len(loading_info['missing_keys'])   ,      0)

        model.to(lowercase)
        a__:					Any      							=     model(**self.dummy_input)[0]

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



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





        '''simple docstring'''





        a__							,   a__:					Tuple      							=     self.prepare_init_args_and_inputs_for_common()

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

        a__:					List[Any]      							=     ['hidden_states', 'timestep']
        self.assertListEqual(arg_names[:2]   ,      lowercase)



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





        '''simple docstring'''





        a__:					str      							=     PriorTransformer.from_pretrained('hf-internal-testing/prior-dummy')
        a__:					str      							=     model.to(lowercase)

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

        a__:					Dict      							=     self.get_dummy_seed_input()

        with torch.no_grad():
          a__:					str      							=     model(**lowercase)[0]

        a__:					str      							=     output[0, :5].flatten().cpu()
        print(lowercase)

        # 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__:					Any      							=     torch.tensor([-1.3436, -0.2870, 0.7538, 0.4368, -0.0239])
        self.assertTrue(torch_all_close(lowercase   ,      lowercase   ,      rtol=1e-2))







@slow
class        __snake_case	(					unittest.TestCase    ):



      def   lowerCamelCase_       (						self   ,      lowercase=1   ,      lowercase=7_68   ,      lowercase=77   ,      lowercase=0)     ->      int:





        '''simple docstring'''





        torch.manual_seed(lowercase)
        a__:					Union[str, Any]      							=     batch_size
        a__:					List[str]      							=     embedding_dim
        a__:					str      							=     num_embeddings

        a__:					Tuple      							=     torch.randn((batch_size, embedding_dim)).to(lowercase)

        a__:					List[str]      							=     torch.randn((batch_size, embedding_dim)).to(lowercase)
        a__:					str      							=     torch.randn((batch_size, num_embeddings, embedding_dim)).to(lowercase)

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



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





        '''simple docstring'''





        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   lowerCamelCase_       (						self   ,      lowercase   ,      lowercase)     ->      str:





        '''simple docstring'''





        a__:					Tuple      							=     PriorTransformer.from_pretrained('kandinsky-community/kandinsky-2-1-prior'   ,      subfolder='prior')
        model.to(lowercase)
        a__:					Optional[Any]      							=     self.get_dummy_seed_input(seed=lowercase)

        with torch.no_grad():
          a__:					Optional[int]      							=     model(**lowercase)[0]

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

        a__:					List[str]      							=     sample[0, :8].flatten().cpu()
        print(lowercase)
        a__:					Union[str, Any]      							=     torch.tensor(lowercase)

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


