


'''simple docstring'''






import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class 						a__							(				a__							,							unittest.TestCase    ):





							'''simple docstring'''






							lowercase__     :  Optional[Any]      		= KandinskyVaaPriorPipeline
							lowercase__     :  List[str]      		= ["prompt"]
							lowercase__     :  Union[str, Any]      		= ["prompt", "negative_prompt"]
							lowercase__     :  Optional[int]      		= [
							    "num_images_per_prompt",
							    "generator",
							    "num_inference_steps",
							    "latents",
							    "negative_prompt",
							    "guidance_scale",
							    "output_type",
							    "return_dict",
							]
							lowercase__     :  str      		= False







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Dict:
														return 32







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														return 32







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														return self.time_input_dim







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Union[str, Any]:
														return self.time_input_dim * 4







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														return 1_00







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Union[str, Any]:
														lowerCAmelCase__			  =      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )
														return tokenizer







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      CLIPTextConfig(
														    bos_token_id=0   ,       eos_token_id=2   ,       hidden_size=self.text_embedder_hidden_size   ,       projection_dim=self.text_embedder_hidden_size   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       pad_token_id=1   ,       vocab_size=10_00   ,       )
														return CLIPTextModelWithProjection(lowerCamelCase_  )







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														torch.manual_seed(0  )

														lowerCAmelCase__			  =      {
														    '''num_attention_heads''': 2,
														    '''attention_head_dim''': 12,
														    '''embedding_dim''': self.text_embedder_hidden_size,
														    '''num_layers''': 1,
														}

														lowerCAmelCase__			  =      PriorTransformer(**lowerCamelCase_  )
														# clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
														lowerCAmelCase__			  =      nn.Parameter(torch.ones(model.clip_std.shape  )  )
														return model







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														torch.manual_seed(0  )
														lowerCAmelCase__			  =      CLIPVisionConfig(
														    hidden_size=self.text_embedder_hidden_size   ,       image_size=2_24   ,       projection_dim=self.text_embedder_hidden_size   ,       intermediate_size=37   ,       num_attention_heads=4   ,       num_channels=3   ,       num_hidden_layers=5   ,       patch_size=14   ,       )

														lowerCAmelCase__			  =      CLIPVisionModelWithProjection(lowerCamelCase_  )
														return model







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Any:
														lowerCAmelCase__			  =      CLIPImageProcessor(
														    crop_size=2_24   ,       do_center_crop=lowerCamelCase_   ,       do_normalize=lowerCamelCase_   ,       do_resize=lowerCamelCase_   ,       image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]   ,       image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]   ,       resample=3   ,       size=2_24   ,       )

														return image_processor







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Dict:
														lowerCAmelCase__			  =      self.dummy_prior
														lowerCAmelCase__			  =      self.dummy_image_encoder
														lowerCAmelCase__			  =      self.dummy_text_encoder
														lowerCAmelCase__			  =      self.dummy_tokenizer
														lowerCAmelCase__			  =      self.dummy_image_processor

														lowerCAmelCase__			  =      UnCLIPScheduler(
														    variance_type='''fixed_small_log'''   ,       prediction_type='''sample'''   ,       num_train_timesteps=10_00   ,       clip_sample=lowerCamelCase_   ,       clip_sample_range=10.0   ,       )

														lowerCAmelCase__			  =      {
														    '''prior''': prior,
														    '''image_encoder''': image_encoder,
														    '''text_encoder''': text_encoder,
														    '''tokenizer''': tokenizer,
														    '''scheduler''': scheduler,
														    '''image_processor''': image_processor,
														}

														return components







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_=0  )			->				int:
														if str(lowerCamelCase_  ).startswith('''mps'''  ):
																					lowerCAmelCase__			  =      torch.manual_seed(lowerCamelCase_  )
														else:
																					lowerCAmelCase__			  =      torch.Generator(device=lowerCamelCase_  ).manual_seed(lowerCamelCase_  )
														lowerCAmelCase__			  =      {
														    '''prompt''': '''horse''',
														    '''generator''': generator,
														    '''guidance_scale''': 4.0,
														    '''num_inference_steps''': 2,
														    '''output_type''': '''np''',
														}
														return inputs







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														lowerCAmelCase__			  =      '''cpu'''

														lowerCAmelCase__			  =      self.get_dummy_components()

														lowerCAmelCase__			  =      self.pipeline_class(**lowerCamelCase_  )
														lowerCAmelCase__			  =      pipe.to(lowerCamelCase_  )

														pipe.set_progress_bar_config(disable=lowerCamelCase_  )

														lowerCAmelCase__			  =      pipe(**self.get_dummy_inputs(lowerCamelCase_  )  )
														lowerCAmelCase__			  =      output.image_embeds

														lowerCAmelCase__			  =      pipe(
														    **self.get_dummy_inputs(lowerCamelCase_  )   ,       return_dict=lowerCamelCase_   ,       )[0]

														lowerCAmelCase__			  =      image[0, -10:]
														lowerCAmelCase__			  =      image_from_tuple[0, -10:]

														assert image.shape == (1, 32)

														lowerCAmelCase__			  =      np.array(
														    [-0.0_532, 1.7_120, 0.3_656, -1.0_852, -0.8_946, -1.1_756, 0.4_348, 0.2_482, 0.5_146, -0.1_156]  )

														assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2
														assert np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1e-2







							@skip_mps
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														lowerCAmelCase__			  =      torch_device == '''cpu'''
														lowerCAmelCase__			  =      True
														lowerCAmelCase__			  =      False

														self._test_inference_batch_single_identical(
														    test_max_difference=lowerCamelCase_   ,       relax_max_difference=lowerCamelCase_   ,       test_mean_pixel_difference=lowerCamelCase_   ,       )







							@skip_mps
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														lowerCAmelCase__			  =      torch_device == '''cpu'''
														lowerCAmelCase__			  =      False

														self._test_attention_slicing_forward_pass(
														    test_max_difference=lowerCamelCase_   ,       test_mean_pixel_difference=lowerCamelCase_   ,       )



'''simple docstring'''






import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    BertTokenizer,
    ViltConfig,
    ViltForImageAndTextRetrieval,
    ViltForImagesAndTextClassification,
    ViltForMaskedLM,
    ViltForQuestionAnswering,
    ViltImageProcessor,
    ViltProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCAmelCase       							=				logging.get_logger(__name__)


def    _snake_case  (	A					,       A=False					,       A=False					,       A=False			)		->       Union[str, Any]:
							lowerCAmelCase__			  =      []
							for i in range(config.num_hidden_layers			):
														# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
														rename_keys.append((F"""transformer.blocks.{i}.norm1.weight""", F"""vilt.encoder.layer.{i}.layernorm_before.weight""")			)
														rename_keys.append((F"""transformer.blocks.{i}.norm1.bias""", F"""vilt.encoder.layer.{i}.layernorm_before.bias""")			)
														rename_keys.append(
														    (F"""transformer.blocks.{i}.attn.proj.weight""", F"""vilt.encoder.layer.{i}.attention.output.dense.weight""")			)
														rename_keys.append(
														    (F"""transformer.blocks.{i}.attn.proj.bias""", F"""vilt.encoder.layer.{i}.attention.output.dense.bias""")			)
														rename_keys.append((F"""transformer.blocks.{i}.norm2.weight""", F"""vilt.encoder.layer.{i}.layernorm_after.weight""")			)
														rename_keys.append((F"""transformer.blocks.{i}.norm2.bias""", F"""vilt.encoder.layer.{i}.layernorm_after.bias""")			)
														rename_keys.append(
														    (F"""transformer.blocks.{i}.mlp.fc1.weight""", F"""vilt.encoder.layer.{i}.intermediate.dense.weight""")			)
														rename_keys.append((F"""transformer.blocks.{i}.mlp.fc1.bias""", F"""vilt.encoder.layer.{i}.intermediate.dense.bias""")			)
														rename_keys.append((F"""transformer.blocks.{i}.mlp.fc2.weight""", F"""vilt.encoder.layer.{i}.output.dense.weight""")			)
														rename_keys.append((F"""transformer.blocks.{i}.mlp.fc2.bias""", F"""vilt.encoder.layer.{i}.output.dense.bias""")			)

							# embeddings
							rename_keys.extend(
							    [
							        # text embeddings
							        ('''text_embeddings.word_embeddings.weight''', '''vilt.embeddings.text_embeddings.word_embeddings.weight'''),
							        (
							            '''text_embeddings.position_embeddings.weight''',
							            '''vilt.embeddings.text_embeddings.position_embeddings.weight''',
							        ),
							        ('''text_embeddings.position_ids''', '''vilt.embeddings.text_embeddings.position_ids'''),
							        (
							            '''text_embeddings.token_type_embeddings.weight''',
							            '''vilt.embeddings.text_embeddings.token_type_embeddings.weight''',
							        ),
							        ('''text_embeddings.LayerNorm.weight''', '''vilt.embeddings.text_embeddings.LayerNorm.weight'''),
							        ('''text_embeddings.LayerNorm.bias''', '''vilt.embeddings.text_embeddings.LayerNorm.bias'''),
							        # patch embeddings
							        ('''transformer.cls_token''', '''vilt.embeddings.cls_token'''),
							        ('''transformer.patch_embed.proj.weight''', '''vilt.embeddings.patch_embeddings.projection.weight'''),
							        ('''transformer.patch_embed.proj.bias''', '''vilt.embeddings.patch_embeddings.projection.bias'''),
							        ('''transformer.pos_embed''', '''vilt.embeddings.position_embeddings'''),
							        # token type embeddings
							        ('''token_type_embeddings.weight''', '''vilt.embeddings.token_type_embeddings.weight'''),
							    ]			)

							# final layernorm + pooler
							rename_keys.extend(
							    [
							        ('''transformer.norm.weight''', '''vilt.layernorm.weight'''),
							        ('''transformer.norm.bias''', '''vilt.layernorm.bias'''),
							        ('''pooler.dense.weight''', '''vilt.pooler.dense.weight'''),
							        ('''pooler.dense.bias''', '''vilt.pooler.dense.bias'''),
							    ]			)

							# classifier head(s)
							if vqa_model:
														# classification head
														rename_keys.extend(
														    [
														        ('''vqa_classifier.0.weight''', '''classifier.0.weight'''),
														        ('''vqa_classifier.0.bias''', '''classifier.0.bias'''),
														        ('''vqa_classifier.1.weight''', '''classifier.1.weight'''),
														        ('''vqa_classifier.1.bias''', '''classifier.1.bias'''),
														        ('''vqa_classifier.3.weight''', '''classifier.3.weight'''),
														        ('''vqa_classifier.3.bias''', '''classifier.3.bias'''),
														    ]			)
							elif nlvr_model:
														# classification head
														rename_keys.extend(
														    [
														        ('''nlvr2_classifier.0.weight''', '''classifier.0.weight'''),
														        ('''nlvr2_classifier.0.bias''', '''classifier.0.bias'''),
														        ('''nlvr2_classifier.1.weight''', '''classifier.1.weight'''),
														        ('''nlvr2_classifier.1.bias''', '''classifier.1.bias'''),
														        ('''nlvr2_classifier.3.weight''', '''classifier.3.weight'''),
														        ('''nlvr2_classifier.3.bias''', '''classifier.3.bias'''),
														    ]			)
							else:
														pass

							return rename_keys


def    _snake_case  (	A					,       A			)		->       List[str]:
							for i in range(config.num_hidden_layers			):
														lowerCAmelCase__			  =      '''vilt.'''
														# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
														lowerCAmelCase__			  =      state_dict.pop(F"""transformer.blocks.{i}.attn.qkv.weight"""			)
														lowerCAmelCase__			  =      state_dict.pop(F"""transformer.blocks.{i}.attn.qkv.bias"""			)
														# next, add query, keys and values (in that order) to the state dict
														lowerCAmelCase__			  =      in_proj_weight[
														    : config.hidden_size, :
														]
														lowerCAmelCase__			  =      in_proj_bias[: config.hidden_size]
														lowerCAmelCase__			  =      in_proj_weight[
														    config.hidden_size : config.hidden_size * 2, :
														]
														lowerCAmelCase__			  =      in_proj_bias[
														    config.hidden_size : config.hidden_size * 2
														]
														lowerCAmelCase__			  =      in_proj_weight[
														    -config.hidden_size :, :
														]
														lowerCAmelCase__			  =      in_proj_bias[-config.hidden_size :]


def    _snake_case  (	A			)		->       List[str]:
							lowerCAmelCase__			  =      ['''head.weight''', '''head.bias''']
							for k in ignore_keys:
														state_dict.pop(A					,       A			)


def    _snake_case  (	A					,       A					,       A			)		->       str:
							lowerCAmelCase__			  =      dct.pop(A			)
							lowerCAmelCase__			  =      val


@torch.no_grad()
def    _snake_case  (	A					,       A			)		->       Any:
							lowerCAmelCase__			  =      ViltConfig(image_size=384					,       patch_size=32					,       tie_word_embeddings=A			)
							lowerCAmelCase__			  =      False
							lowerCAmelCase__			  =      False
							lowerCAmelCase__			  =      False
							lowerCAmelCase__			  =      False
							if "vqa" in checkpoint_url:
														lowerCAmelCase__			  =      True
														lowerCAmelCase__			  =      3129
														lowerCAmelCase__			  =      '''huggingface/label-files'''
														lowerCAmelCase__			  =      '''vqa2-id2label.json'''
														lowerCAmelCase__			  =      json.load(open(hf_hub_download(A					,       A					,       repo_type='''dataset'''			)					,       '''r'''			)			)
														lowerCAmelCase__			  =      {int(A			): v for k, v in idalabel.items()}
														lowerCAmelCase__			  =      idalabel
														lowerCAmelCase__			  =      {v: k for k, v in idalabel.items()}
														lowerCAmelCase__			  =      ViltForQuestionAnswering(A			)
							elif "nlvr" in checkpoint_url:
														lowerCAmelCase__			  =      True
														lowerCAmelCase__			  =      2
														lowerCAmelCase__			  =      {0: '''False''', 1: '''True'''}
														lowerCAmelCase__			  =      {v: k for k, v in config.idalabel.items()}
														lowerCAmelCase__			  =      3
														lowerCAmelCase__			  =      ViltForImagesAndTextClassification(A			)
							elif "irtr" in checkpoint_url:
														lowerCAmelCase__			  =      True
														lowerCAmelCase__			  =      ViltForImageAndTextRetrieval(A			)
							elif "mlm_itm" in checkpoint_url:
														lowerCAmelCase__			  =      True
														lowerCAmelCase__			  =      ViltForMaskedLM(A			)
							else:
														raise ValueError('''Unknown model type'''			)

							# load state_dict of original model, remove and rename some keys
							lowerCAmelCase__			  =      torch.hub.load_state_dict_from_url(A					,       map_location='''cpu'''			)['''state_dict''']
							lowerCAmelCase__			  =      create_rename_keys(A					,       A					,       A					,       A			)
							for src, dest in rename_keys:
														rename_key(A					,       A					,       A			)
							read_in_q_k_v(A					,       A			)
							if mlm_model or irtr_model:
														lowerCAmelCase__			  =      ['''itm_score.fc.weight''', '''itm_score.fc.bias''']
														for k in ignore_keys:
																					state_dict.pop(A					,       A			)

    # load state dict into HuggingFace model
							model.eval()
							if mlm_model:
														lowerCAmelCase__  ,   lowerCAmelCase__			  =      model.load_state_dict(A					,       strict=A			)
														assert missing_keys == ["mlm_score.decoder.bias"]
							else:
														model.load_state_dict(A			)

							# Define processor
							lowerCAmelCase__			  =      ViltImageProcessor(size=384			)
							lowerCAmelCase__			  =      BertTokenizer.from_pretrained('''bert-base-uncased'''			)
							lowerCAmelCase__			  =      ViltProcessor(A					,       A			)

							# Forward pass on example inputs (image + text)
							if nlvr_model:
														lowerCAmelCase__			  =      Image.open(requests.get('''https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg'''					,       stream=A			).raw			)
														lowerCAmelCase__			  =      Image.open(requests.get('''https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg'''					,       stream=A			).raw			)
														lowerCAmelCase__			  =      (
														    '''The left image contains twice the number of dogs as the right image, and at least two dogs in total are'''
														    ''' standing.'''
														)
														lowerCAmelCase__			  =      processor(A					,       A					,       return_tensors='''pt'''			)
														lowerCAmelCase__			  =      processor(A					,       A					,       return_tensors='''pt'''			)
														lowerCAmelCase__			  =      model(
														    input_ids=encoding_a.input_ids					,       pixel_values=encoding_a.pixel_values					,       pixel_values_a=encoding_a.pixel_values					,       )
							else:
														lowerCAmelCase__			  =      Image.open(requests.get('''http://images.cocodataset.org/val2017/000000039769.jpg'''					,       stream=A			).raw			)
														if mlm_model:
																					lowerCAmelCase__			  =      '''a bunch of [MASK] laying on a [MASK].'''
														else:
																					lowerCAmelCase__			  =      '''How many cats are there?'''
														lowerCAmelCase__			  =      processor(A					,       A					,       return_tensors='''pt'''			)
														lowerCAmelCase__			  =      model(**A			)

							# Verify outputs
							if mlm_model:
														lowerCAmelCase__			  =      torch.Size([1, 11, 30522]			)
														lowerCAmelCase__			  =      torch.tensor([-12.5_061, -12.5_123, -12.5_174]			)
														assert outputs.logits.shape == expected_shape
														assert torch.allclose(outputs.logits[0, 0, :3]					,       A					,       atol=1E-4			)

														# verify masked token prediction equals "cats"
														lowerCAmelCase__			  =      outputs.logits[0, 4, :].argmax(-1			).item()
														assert tokenizer.decode([predicted_id]			) == "cats"
							elif vqa_model:
														lowerCAmelCase__			  =      torch.Size([1, 3129]			)
														lowerCAmelCase__			  =      torch.tensor([-15.9_495, -18.1_472, -10.3_041]			)
														assert torch.allclose(outputs.logits[0, :3]					,       A					,       atol=1E-4			)
														assert outputs.logits.shape == expected_shape
														assert torch.allclose(outputs.logits[0, 0, :3]					,       A					,       atol=1E-4			)

														# verify vqa prediction equals "2"
														lowerCAmelCase__			  =      outputs.logits.argmax(-1			).item()
														assert model.config.idalabel[predicted_idx] == "2"
							elif nlvr_model:
														lowerCAmelCase__			  =      torch.Size([1, 2]			)
														lowerCAmelCase__			  =      torch.tensor([-2.8_721, 2.1_291]			)
														assert torch.allclose(outputs.logits[0, :3]					,       A					,       atol=1E-4			)
														assert outputs.logits.shape == expected_shape

							Path(A			).mkdir(exist_ok=A			)
							print(F"""Saving model and processor to {pytorch_dump_folder_path}"""			)
							model.save_pretrained(A			)
							processor.save_pretrained(A			)


if __name__ == "__main__":
		__UpperCAmelCase       							=				argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--checkpoint_url''',
		    default='''https://github.com/dandelin/ViLT/releases/download/200k/vilt_200k_mlm_itm.ckpt''',
		    type=str,
		    help='''URL of the checkpoint you\'d like to convert.''',
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
		)

		__UpperCAmelCase       							=				parser.parse_args()
		convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)