


'''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()
lowerCamelCase__					=   logging.get_logger(__name__)






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase=False					,       __lowerCAmelCase=False					,       __lowerCAmelCase=False      ):
		_UpperCAmelCase						:				Tuple	    =					[]
		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        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase      ):
		for i in range(config.num_hidden_layers      ):
				_UpperCAmelCase						:				Optional[int]	    =					"vilt."
				# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
				_UpperCAmelCase						:				int	    =					state_dict.pop(F"""transformer.blocks.{i}.attn.qkv.weight"""      )
				_UpperCAmelCase						:				Optional[int]	    =					state_dict.pop(F"""transformer.blocks.{i}.attn.qkv.bias"""      )
				# next, add query, keys and values (in that order) to the state dict
				_UpperCAmelCase						:				Optional[Any]	    =					in_proj_weight[
				    : config.hidden_size, :
				]
				_UpperCAmelCase						:				Any	    =					in_proj_bias[: config.hidden_size]
				_UpperCAmelCase						:				Optional[Any]	    =					in_proj_weight[
				    config.hidden_size : config.hidden_size * 2, :
				]
				_UpperCAmelCase						:				str	    =					in_proj_bias[
				    config.hidden_size : config.hidden_size * 2
				]
				_UpperCAmelCase						:				Union[str, Any]	    =					in_proj_weight[
				    -config.hidden_size :, :
				]
				_UpperCAmelCase						:				List[str]	    =					in_proj_bias[-config.hidden_size :]






def        __lowerCAmelCase       (__lowerCAmelCase      ):
		_UpperCAmelCase						:				Any	    =					["head.weight", "head.bias"]
		for k in ignore_keys:
				state_dict.pop(__lowerCAmelCase					,       __lowerCAmelCase      )






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				List[str]	    =					dct.pop(__lowerCAmelCase      )
		_UpperCAmelCase						:				int	    =					val






@torch.no_grad()
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				int	    =					ViltConfig(image_size=384					,       patch_size=32					,       tie_word_embeddings=__lowerCAmelCase      )
		_UpperCAmelCase						:				Optional[Any]	    =					False
		_UpperCAmelCase						:				Optional[Any]	    =					False
		_UpperCAmelCase						:				Tuple	    =					False
		_UpperCAmelCase						:				Dict	    =					False
		if "vqa" in checkpoint_url:
				_UpperCAmelCase						:				Optional[Any]	    =					True
				_UpperCAmelCase						:				Optional[Any]	    =					3_129
				_UpperCAmelCase						:				int	    =					"huggingface/label-files"
				_UpperCAmelCase						:				int	    =					"vqa2-id2label.json"
				_UpperCAmelCase						:				Dict	    =					json.load(open(hf_hub_download(__lowerCAmelCase					,       __lowerCAmelCase					,       repo_type="dataset"      )					,       "r"      )      )
				_UpperCAmelCase						:				Tuple	    =					{int(__lowerCAmelCase      ): v for k, v in idalabel.items()}
				_UpperCAmelCase						:				Optional[Any]	    =					idalabel
				_UpperCAmelCase						:				List[str]	    =					{v: k for k, v in idalabel.items()}
				_UpperCAmelCase						:				int	    =					ViltForQuestionAnswering(__lowerCAmelCase      )
		elif "nlvr" in checkpoint_url:
				_UpperCAmelCase						:				Any	    =					True
				_UpperCAmelCase						:				int	    =					2
				_UpperCAmelCase						:				Union[str, Any]	    =					{0: "False", 1: "True"}
				_UpperCAmelCase						:				Union[str, Any]	    =					{v: k for k, v in config.idalabel.items()}
				_UpperCAmelCase						:				int	    =					3
				_UpperCAmelCase						:				Optional[int]	    =					ViltForImagesAndTextClassification(__lowerCAmelCase      )
		elif "irtr" in checkpoint_url:
				_UpperCAmelCase						:				Optional[Any]	    =					True
				_UpperCAmelCase						:				Union[str, Any]	    =					ViltForImageAndTextRetrieval(__lowerCAmelCase      )
		elif "mlm_itm" in checkpoint_url:
				_UpperCAmelCase						:				Optional[int]	    =					True
				_UpperCAmelCase						:				Optional[Any]	    =					ViltForMaskedLM(__lowerCAmelCase      )
		else:
				raise ValueError("Unknown model type"      )

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

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

		# Define processor
		_UpperCAmelCase						:				Union[str, Any]	    =					ViltImageProcessor(size=384      )
		_UpperCAmelCase						:				Tuple	    =					BertTokenizer.from_pretrained("bert-base-uncased"      )
		_UpperCAmelCase						:				Optional[int]	    =					ViltProcessor(__lowerCAmelCase					,       __lowerCAmelCase      )

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

		# Verify outputs
		if mlm_model:
				_UpperCAmelCase						:				Tuple	    =					torch.Size([1, 11, 30_522]      )
				_UpperCAmelCase						:				Optional[Any]	    =					torch.tensor([-1_2.5_0_6_1, -1_2.5_1_2_3, -1_2.5_1_7_4]      )
				assert outputs.logits.shape == expected_shape
				assert torch.allclose(outputs.logits[0, 0, :3]					,       __lowerCAmelCase					,       atol=1e-4      )

				# verify masked token prediction equals "cats"
				_UpperCAmelCase						:				Union[str, Any]	    =					outputs.logits[0, 4, :].argmax(-1      ).item()
				assert tokenizer.decode([predicted_id]      ) == "cats"
		elif vqa_model:
				_UpperCAmelCase						:				Union[str, Any]	    =					torch.Size([1, 3_129]      )
				_UpperCAmelCase						:				Optional[int]	    =					torch.tensor([-1_5.9_4_9_5, -1_8.1_4_7_2, -1_0.3_0_4_1]      )
				assert torch.allclose(outputs.logits[0, :3]					,       __lowerCAmelCase					,       atol=1e-4      )
				assert outputs.logits.shape == expected_shape
				assert torch.allclose(outputs.logits[0, 0, :3]					,       __lowerCAmelCase					,       atol=1e-4      )

				# verify vqa prediction equals "2"
				_UpperCAmelCase						:				Any	    =					outputs.logits.argmax(-1      ).item()
				assert model.config.idalabel[predicted_idx] == "2"
		elif nlvr_model:
				_UpperCAmelCase						:				List[str]	    =					torch.Size([1, 2]      )
				_UpperCAmelCase						:				Optional[Any]	    =					torch.tensor([-2.8_7_2_1, 2.1_2_9_1]      )
				assert torch.allclose(outputs.logits[0, :3]					,       __lowerCAmelCase					,       atol=1e-4      )
				assert outputs.logits.shape == expected_shape

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


if __name__ == "__main__":
						lowerCamelCase__					=   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.'
						)

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







'''simple docstring'''


import warnings

from diffusers import StableDiffusionImgaImgPipeline  # noqa F401


warnings.warn(
    'The `image_to_image.py` script is outdated. Please use directly `from diffusers import'
    ' StableDiffusionImg2ImgPipeline` instead.'
)



