





import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


# =================#
# UNet Conversion #
# =================#

UpperCAmelCase_  							=				[
    # (stable-diffusion, HF Diffusers)
    ('time_embed.0.weight', 'time_embedding.linear_1.weight'),
    ('time_embed.0.bias', 'time_embedding.linear_1.bias'),
    ('time_embed.2.weight', 'time_embedding.linear_2.weight'),
    ('time_embed.2.bias', 'time_embedding.linear_2.bias'),
    ('input_blocks.0.0.weight', 'conv_in.weight'),
    ('input_blocks.0.0.bias', 'conv_in.bias'),
    ('out.0.weight', 'conv_norm_out.weight'),
    ('out.0.bias', 'conv_norm_out.bias'),
    ('out.2.weight', 'conv_out.weight'),
    ('out.2.bias', 'conv_out.bias'),
]

UpperCAmelCase_  							=				[
    # (stable-diffusion, HF Diffusers)
    ('in_layers.0', 'norm1'),
    ('in_layers.2', 'conv1'),
    ('out_layers.0', 'norm2'),
    ('out_layers.3', 'conv2'),
    ('emb_layers.1', 'time_emb_proj'),
    ('skip_connection', 'conv_shortcut'),
]

UpperCAmelCase_  							=				[]
# hardcoded number of downblocks and resnets/attentions...
# would need smarter logic for other networks.
for i in range(4):
  # loop over downblocks/upblocks

  for j in range(2):
    # loop over resnets/attentions for downblocks
    UpperCAmelCase_  							=				F'''down_blocks.{i}.resnets.{j}.'''
    UpperCAmelCase_  							=				F'''input_blocks.{3*i + j + 1}.0.'''
    unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))

    if i < 3:
      # no attention layers in down_blocks.3
      UpperCAmelCase_  							=				F'''down_blocks.{i}.attentions.{j}.'''
      UpperCAmelCase_  							=				F'''input_blocks.{3*i + j + 1}.1.'''
      unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))

  for j in range(3):
    # loop over resnets/attentions for upblocks
    UpperCAmelCase_  							=				F'''up_blocks.{i}.resnets.{j}.'''
    UpperCAmelCase_  							=				F'''output_blocks.{3*i + j}.0.'''
    unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))

    if i > 0:
      # no attention layers in up_blocks.0
      UpperCAmelCase_  							=				F'''up_blocks.{i}.attentions.{j}.'''
      UpperCAmelCase_  							=				F'''output_blocks.{3*i + j}.1.'''
      unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))

  if i < 3:
    # no downsample in down_blocks.3
    UpperCAmelCase_  							=				F'''down_blocks.{i}.downsamplers.0.conv.'''
    UpperCAmelCase_  							=				F'''input_blocks.{3*(i+1)}.0.op.'''
    unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))

    # no upsample in up_blocks.3
    UpperCAmelCase_  							=				F'''up_blocks.{i}.upsamplers.0.'''
    UpperCAmelCase_  							=				F'''output_blocks.{3*i + 2}.{1 if i == 0 else 2}.'''
    unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

UpperCAmelCase_  							=				'mid_block.attentions.0.'
UpperCAmelCase_  							=				'middle_block.1.'
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
  UpperCAmelCase_  							=				F'''mid_block.resnets.{j}.'''
  UpperCAmelCase_  							=				F'''middle_block.{2*j}.'''
  unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))




def 	lowerCAmelCase_					(						__UpperCAmelCase:    int   )   ->			str:
    # buyer beware: this is a *brittle* function,
    # and correct output requires that all of these pieces interact in
    # the exact order in which I have arranged them.
    UpperCamelCase__							:       List[Any]      =   {k: k for k in unet_state_dict.keys()}
    for sd_name, hf_name in unet_conversion_map:
        UpperCamelCase__							:       Tuple      =   sd_name
    for k, v in mapping.items():
        if "resnets" in k:
            for sd_part, hf_part in unet_conversion_map_resnet:
                UpperCamelCase__							:       Optional[int]      =   v.replace(__UpperCAmelCase	,     __UpperCAmelCase   )
            UpperCamelCase__							:       str      =   v
    for k, v in mapping.items():
        for sd_part, hf_part in unet_conversion_map_layer:
            UpperCamelCase__							:       List[Any]      =   v.replace(__UpperCAmelCase	,     __UpperCAmelCase   )
        UpperCamelCase__							:       List[str]      =   v
    UpperCamelCase__							:       Tuple      =   {v: unet_state_dict[k] for k, v in mapping.items()}
    return new_state_dict


# ================#
# VAE Conversion #
# ================#

UpperCAmelCase_  							=				[
    # (stable-diffusion, HF Diffusers)
    ('nin_shortcut', 'conv_shortcut'),
    ('norm_out', 'conv_norm_out'),
    ('mid.attn_1.', 'mid_block.attentions.0.'),
]

for i in range(4):
  # down_blocks have two resnets
  for j in range(2):
    UpperCAmelCase_  							=				F'''encoder.down_blocks.{i}.resnets.{j}.'''
    UpperCAmelCase_  							=				F'''encoder.down.{i}.block.{j}.'''
    vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

  if i < 3:
    UpperCAmelCase_  							=				F'''down_blocks.{i}.downsamplers.0.'''
    UpperCAmelCase_  							=				F'''down.{i}.downsample.'''
    vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

    UpperCAmelCase_  							=				F'''up_blocks.{i}.upsamplers.0.'''
    UpperCAmelCase_  							=				F'''up.{3-i}.upsample.'''
    vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))

  # up_blocks have three resnets
  # also, up blocks in hf are numbered in reverse from sd
  for j in range(3):
    UpperCAmelCase_  							=				F'''decoder.up_blocks.{i}.resnets.{j}.'''
    UpperCAmelCase_  							=				F'''decoder.up.{3-i}.block.{j}.'''
    vae_conversion_map.append((sd_up_prefix, hf_up_prefix))

# this part accounts for mid blocks in both the encoder and the decoder
for i in range(2):
  UpperCAmelCase_  							=				F'''mid_block.resnets.{i}.'''
  UpperCAmelCase_  							=				F'''mid.block_{i+1}.'''
  vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


UpperCAmelCase_  							=				[
    # (stable-diffusion, HF Diffusers)
    ('norm.', 'group_norm.'),
    ('q.', 'query.'),
    ('k.', 'key.'),
    ('v.', 'value.'),
    ('proj_out.', 'proj_attn.'),
]




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Optional[int]   )   ->			List[Any]:
    # convert HF linear weights to SD conv2d weights
    return w.reshape(*w.shape	,     1	,     1   )




def 	lowerCAmelCase_					(						__UpperCAmelCase:    int   )   ->			List[Any]:
    UpperCamelCase__							:       List[str]      =   {k: k for k in vae_state_dict.keys()}
    for k, v in mapping.items():
        for sd_part, hf_part in vae_conversion_map:
            UpperCamelCase__							:       List[str]      =   v.replace(__UpperCAmelCase	,     __UpperCAmelCase   )
        UpperCamelCase__							:       Optional[int]      =   v
    for k, v in mapping.items():
        if "attentions" in k:
            for sd_part, hf_part in vae_conversion_map_attn:
                UpperCamelCase__							:       Any      =   v.replace(__UpperCAmelCase	,     __UpperCAmelCase   )
            UpperCamelCase__							:       str      =   v
    UpperCamelCase__							:       Any      =   {v: vae_state_dict[k] for k, v in mapping.items()}
    UpperCamelCase__							:       Any      =   ['''q''', '''k''', '''v''', '''proj_out''']
    for k, v in new_state_dict.items():
        for weight_name in weights_to_convert:
            if f"mid.attn_1.{weight_name}.weight" in k:
                print(f"Reshaping {k} for SD format"   )
                UpperCamelCase__							:       Dict      =   reshape_weight_for_sd(__UpperCAmelCase   )
    return new_state_dict


# =========================#
# Text Encoder Conversion #
# =========================#


UpperCAmelCase_  							=				[
    # (stable-diffusion, HF Diffusers)
    ('resblocks.', 'text_model.encoder.layers.'),
    ('ln_1', 'layer_norm1'),
    ('ln_2', 'layer_norm2'),
    ('.c_fc.', '.fc1.'),
    ('.c_proj.', '.fc2.'),
    ('.attn', '.self_attn'),
    ('ln_final.', 'transformer.text_model.final_layer_norm.'),
    ('token_embedding.weight', 'transformer.text_model.embeddings.token_embedding.weight'),
    ('positional_embedding', 'transformer.text_model.embeddings.position_embedding.weight'),
]
UpperCAmelCase_  							=				{re.escape(x[1]): x[0] for x in textenc_conversion_lst}
UpperCAmelCase_  							=				re.compile('|'.join(protected.keys()))

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
UpperCAmelCase_  							=				{'q': 0, 'k': 1, 'v': 2}




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Union[str, Any]   )   ->			Dict:
    UpperCamelCase__							:       Tuple      =   {}
    UpperCamelCase__							:       Any      =   {}
    UpperCamelCase__							:       Dict      =   {}
    for k, v in text_enc_dict.items():
        if (
            k.endswith('''.self_attn.q_proj.weight'''   )
            or k.endswith('''.self_attn.k_proj.weight'''   )
            or k.endswith('''.self_attn.v_proj.weight'''   )
        ):
            UpperCamelCase__							:       Union[str, Any]      =   k[: -len('''.q_proj.weight'''   )]
            UpperCamelCase__							:       Tuple      =   k[-len('''q_proj.weight'''   )]
            if k_pre not in capture_qkv_weight:
                UpperCamelCase__							:       str      =   [None, None, None]
            UpperCamelCase__							:       List[str]      =   v
            continue

        if (
            k.endswith('''.self_attn.q_proj.bias'''   )
            or k.endswith('''.self_attn.k_proj.bias'''   )
            or k.endswith('''.self_attn.v_proj.bias'''   )
        ):
            UpperCamelCase__							:       Optional[Any]      =   k[: -len('''.q_proj.bias'''   )]
            UpperCamelCase__							:       Union[str, Any]      =   k[-len('''q_proj.bias'''   )]
            if k_pre not in capture_qkv_bias:
                UpperCamelCase__							:       Any      =   [None, None, None]
            UpperCamelCase__							:       Dict      =   v
            continue

        UpperCamelCase__							:       List[Any]      =   textenc_pattern.sub(lambda __UpperCAmelCase   : protected[re.escape(m.group(0   )   )]	,     __UpperCAmelCase   )
        UpperCamelCase__							:       Dict      =   v

    for k_pre, tensors in capture_qkv_weight.items():
        if None in tensors:
            raise Exception('''CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing'''   )
        UpperCamelCase__							:       Optional[Any]      =   textenc_pattern.sub(lambda __UpperCAmelCase   : protected[re.escape(m.group(0   )   )]	,     __UpperCAmelCase   )
        UpperCamelCase__							:       str      =   torch.cat(__UpperCAmelCase   )

    for k_pre, tensors in capture_qkv_bias.items():
        if None in tensors:
            raise Exception('''CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing'''   )
        UpperCamelCase__							:       List[str]      =   textenc_pattern.sub(lambda __UpperCAmelCase   : protected[re.escape(m.group(0   )   )]	,     __UpperCAmelCase   )
        UpperCamelCase__							:       str      =   torch.cat(__UpperCAmelCase   )

    return new_state_dict




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Union[str, Any]   )   ->			List[str]:
    return text_enc_dict


if __name__ == "__main__":
  UpperCAmelCase_  							=				argparse.ArgumentParser()

  parser.add_argument('--model_path', default=None, type=str, required=True, help='Path to the model to convert.')
  parser.add_argument('--checkpoint_path', default=None, type=str, required=True, help='Path to the output model.')
  parser.add_argument('--half', action='store_true', help='Save weights in half precision.')
  parser.add_argument(
      '--use_safetensors', action='store_true', help='Save weights use safetensors, default is ckpt.'
  )

  UpperCAmelCase_  							=				parser.parse_args()

  assert args.model_path is not None, "Must provide a model path!"

  assert args.checkpoint_path is not None, "Must provide a checkpoint path!"

  # Path for safetensors
  UpperCAmelCase_  							=				osp.join(args.model_path, 'unet', 'diffusion_pytorch_model.safetensors')
  UpperCAmelCase_  							=				osp.join(args.model_path, 'vae', 'diffusion_pytorch_model.safetensors')
  UpperCAmelCase_  							=				osp.join(args.model_path, 'text_encoder', 'model.safetensors')

  # Load models from safetensors if it exists, if it doesn't pytorch
  if osp.exists(unet_path):
    UpperCAmelCase_  							=				load_file(unet_path, device='cpu')
  else:
    UpperCAmelCase_  							=				osp.join(args.model_path, 'unet', 'diffusion_pytorch_model.bin')
    UpperCAmelCase_  							=				torch.load(unet_path, map_location='cpu')

  if osp.exists(vae_path):
    UpperCAmelCase_  							=				load_file(vae_path, device='cpu')
  else:
    UpperCAmelCase_  							=				osp.join(args.model_path, 'vae', 'diffusion_pytorch_model.bin')
    UpperCAmelCase_  							=				torch.load(vae_path, map_location='cpu')

  if osp.exists(text_enc_path):
    UpperCAmelCase_  							=				load_file(text_enc_path, device='cpu')
  else:
    UpperCAmelCase_  							=				osp.join(args.model_path, 'text_encoder', 'pytorch_model.bin')
    UpperCAmelCase_  							=				torch.load(text_enc_path, map_location='cpu')

  # Convert the UNet model
  UpperCAmelCase_  							=				convert_unet_state_dict(unet_state_dict)
  UpperCAmelCase_  							=				{'model.diffusion_model.' + k: v for k, v in unet_state_dict.items()}

  # Convert the VAE model
  UpperCAmelCase_  							=				convert_vae_state_dict(vae_state_dict)
  UpperCAmelCase_  							=				{'first_stage_model.' + k: v for k, v in vae_state_dict.items()}

  # Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
  UpperCAmelCase_  							=				'text_model.encoder.layers.22.layer_norm2.bias' in text_enc_dict

  if is_vaa_model:
    # Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
    UpperCAmelCase_  							=				{'transformer.' + k: v for k, v in text_enc_dict.items()}
    UpperCAmelCase_  							=				convert_text_enc_state_dict_vaa(text_enc_dict)
    UpperCAmelCase_  							=				{'cond_stage_model.model.' + k: v for k, v in text_enc_dict.items()}
  else:
    UpperCAmelCase_  							=				convert_text_enc_state_dict(text_enc_dict)
    UpperCAmelCase_  							=				{'cond_stage_model.transformer.' + k: v for k, v in text_enc_dict.items()}

  # Put together new checkpoint
  UpperCAmelCase_  							=				{**unet_state_dict, **vae_state_dict, **text_enc_dict}
  if args.half:
    UpperCAmelCase_  							=				{k: v.half() for k, v in state_dict.items()}

  if args.use_safetensors:
    save_file(state_dict, args.checkpoint_path)
  else:
    UpperCAmelCase_  							=				{'state_dict': state_dict}
    torch.save(state_dict, args.checkpoint_path)








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




def 	lowerCAmelCase_					(						__UpperCAmelCase:    List[Any]   )   ->			Any:
    UpperCamelCase__							:       int      =   ['''head.weight''', '''head.bias''']
    for k in ignore_keys:
        state_dict.pop(__UpperCAmelCase	,     __UpperCAmelCase   )




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Any	,     __UpperCAmelCase:    Dict	,     __UpperCAmelCase:    int   )   ->			List[str]:
    UpperCamelCase__							:       str      =   dct.pop(__UpperCAmelCase   )
    UpperCamelCase__							:       Optional[int]      =   val




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

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

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

    # Define processor
    UpperCamelCase__							:       Union[str, Any]      =   ViltImageProcessor(size=384   )
    UpperCamelCase__							:       List[str]      =   BertTokenizer.from_pretrained('''bert-base-uncased'''   )
    UpperCamelCase__							:       Optional[int]      =   ViltProcessor(__UpperCAmelCase	,     __UpperCAmelCase   )

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

    # Verify outputs
    if mlm_model:
        UpperCamelCase__							:       str      =   torch.Size([1, 11, 3_0522]   )
        UpperCamelCase__							:       Optional[Any]      =   torch.tensor([-12.5061, -12.5123, -12.5174]   )
        assert outputs.logits.shape == expected_shape
        assert torch.allclose(outputs.logits[0, 0, :3]	,     __UpperCAmelCase	,     atol=1e-4   )

        # verify masked token prediction equals "cats"
        UpperCamelCase__							:       Optional[Any]      =   outputs.logits[0, 4, :].argmax(-1   ).item()
        assert tokenizer.decode([predicted_id]   ) == "cats"
    elif vqa_model:
        UpperCamelCase__							:       List[Any]      =   torch.Size([1, 3129]   )
        UpperCamelCase__							:       Tuple      =   torch.tensor([-15.9495, -18.1472, -10.3041]   )
        assert torch.allclose(outputs.logits[0, :3]	,     __UpperCAmelCase	,     atol=1e-4   )
        assert outputs.logits.shape == expected_shape
        assert torch.allclose(outputs.logits[0, 0, :3]	,     __UpperCAmelCase	,     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__							:       Dict      =   torch.Size([1, 2]   )
        UpperCamelCase__							:       str      =   torch.tensor([-2.8721, 2.1291]   )
        assert torch.allclose(outputs.logits[0, :3]	,     __UpperCAmelCase	,     atol=1e-4   )
        assert outputs.logits.shape == expected_shape

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


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)

